diff --git a/crates/vendor/candle/candle-examples/Cargo.toml b/crates/vendor/candle/candle-examples/Cargo.toml
index e679d01b60..4ee6e0b876 100644
--- a/crates/vendor/candle/candle-examples/Cargo.toml
+++ b/crates/vendor/candle/candle-examples/Cargo.toml
@@ -122,3 +122,7 @@ required-features = ["onnx"]
 [[example]]
 name = "colpali"
 required-features = ["pdf2image"]
+
+[[example]]
+name = "lcm_debug_test"
+path = "examples/lcm/debug_test.rs"
\ No newline at end of file
diff --git a/crates/vendor/candle/candle-examples/examples/lcm/main.rs b/crates/vendor/candle/candle-examples/examples/lcm/main.rs
new file mode 100644
index 0000000000..7319f7d584
--- /dev/null
+++ b/crates/vendor/candle/candle-examples/examples/lcm/main.rs
@@ -0,0 +1,872 @@
+#[cfg(feature = "accelerate")]
+extern crate accelerate_src;
+
+#[cfg(feature = "mkl")]
+extern crate intel_mkl_src;
+
+use candle_transformers::models::stable_diffusion;
+use std::ops::Div;
+
+use anyhow::{Error as E, Result};
+use candle::{DType, Device, IndexOp, Module, Tensor, D};
+use clap::Parser;
+use rand::Rng;
+use stable_diffusion::vae::AutoEncoderKL;
+use tokenizers::Tokenizer;
+use candle_nn::VarBuilder;
+use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
+use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModelConfig;
+use candle_transformers::models::stable_diffusion::unet_2d::BlockConfig;
+
+const PRINT_DEBUG_TENSORS: bool = false;
+
+#[derive(Parser)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+  /// The prompt to be used for image generation.
+  #[arg(long, default_value = "A very realistic photo of a rusty robot walking on a sandy beach")]
+  prompt: String,
+
+  #[arg(long, default_value = "")]
+  uncond_prompt: String,
+
+  /// Run on CPU rather than on GPU.
+  #[arg(long)]
+  cpu: 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 width in pixels of the generated image.
+  #[arg(long)]
+  width: Option<usize>,
+
+  /// The UNet weight file, in .safetensors format.
+  #[arg(long, value_name = "FILE")]
+  unet_weights: Option<String>,
+
+  /// The CLIP weight file, in .safetensors format.
+  #[arg(long, value_name = "FILE")]
+  clip_weights: Option<String>,
+
+  /// The CLIP2 weight file, in .safetensors format.
+  #[arg(long, value_name = "FILE")]
+  clip2_weights: Option<String>,
+
+  /// The VAE weight file, in .safetensors format.
+  #[arg(long, value_name = "FILE")]
+  vae_weights: Option<String>,
+
+  #[arg(long, value_name = "FILE")]
+  /// The file specifying the tokenizer to used for tokenization.
+  tokenizer: Option<String>,
+
+  /// The size of the sliced attention or 0 for automatic slicing (disabled by default)
+  #[arg(long)]
+  sliced_attention_size: Option<usize>,
+
+  /// The number of inference steps to run.
+  #[arg(long, default_value_t = 4)]
+  n_steps: usize,
+
+  /// The number of samples to generate iteratively.
+  #[arg(long, default_value_t = 1)]
+  num_samples: usize,
+
+  /// The numbers of samples to generate simultaneously.
+  #[arg[long, default_value_t = 1]]
+  bsize: usize,
+
+  /// The name of the final image to generate.
+  #[arg(long, value_name = "FILE", default_value = "sd_final.png")]
+  final_image: String,
+
+  #[arg(long, value_enum, default_value = "v2-1")]
+  sd_version: StableDiffusionVersion,
+
+  /// Generate intermediary images at each step.
+  #[arg(long, action)]
+  intermediary_images: bool,
+
+  #[arg(long)]
+  use_flash_attn: bool,
+
+  #[arg(long)]
+  use_f16: bool,
+
+  /// The guidance scale for classifier-free guidance.
+  #[arg(long, default_value_t = 2.0)]
+  guidance_scale: f64,
+
+  /// Path to the mask image for inpainting.
+  #[arg(long, value_name = "FILE")]
+  mask_path: Option<String>,
+
+  /// Path to the image used to initialize the latents. For inpainting, this is the image to be masked.
+  #[arg(long, value_name = "FILE")]
+  img2img: Option<String>,
+
+  /// The strength parameter for img2img, between 0 and 1.
+  #[arg(long, default_value_t = 0.7)]
+  strength: f64,
+
+  /// The seed to use when generating random samples.
+  #[arg(long)]
+  seed: Option<u64>,
+
+  /// Force the saved image to update only the masked region
+  #[arg(long)]
+  only_update_masked: bool,
+
+  #[arg(long)]
+  save_tensors: bool,
+
+  /// Use half precision (F16) rather than F32
+  #[arg(long)]
+  half_precision: bool,
+}
+
+#[derive(Debug, Clone, Copy, clap::ValueEnum, PartialEq, Eq)]
+enum StableDiffusionVersion {
+  V1_5,
+  V1_5Inpaint,
+  V2_1,
+  V2Inpaint,
+  Xl,
+  XlInpaint,
+  Turbo,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum ModelFile {
+  Tokenizer,
+  Tokenizer2,
+  Clip,
+  Clip2,
+  Unet,
+  UnetLcm,
+  Vae,
+}
+
+impl StableDiffusionVersion {
+  fn repo(&self) -> &'static str {
+    match self {
+      Self::XlInpaint => "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+      Self::Xl => "stabilityai/stable-diffusion-xl-base-1.0",
+      Self::V2Inpaint => "stabilityai/stable-diffusion-2-inpainting",
+      Self::V2_1 => "stabilityai/stable-diffusion-2-1",
+      Self::V1_5 => "runwayml/stable-diffusion-v1-5",
+      Self::V1_5Inpaint => "stable-diffusion-v1-5/stable-diffusion-inpainting",
+      Self::Turbo => "stabilityai/sdxl-turbo",
+    }
+  }
+
+  fn unet_file(&self, use_f16: bool) -> &'static str {
+    match self {
+      Self::V1_5 | Self::V1_5Inpaint | Self::V2_1 | Self::V2Inpaint | Self::Xl | Self::XlInpaint | Self::Turbo => {
+        if use_f16 {
+          "unet/diffusion_pytorch_model.fp16.safetensors"
+        } else {
+          "unet/diffusion_pytorch_model.safetensors"
+        }
+      },
+    }
+  }
+
+  fn vae_file(&self, use_f16: bool) -> &'static str {
+    match self {
+      Self::V1_5 | Self::V1_5Inpaint | Self::V2_1 | Self::V2Inpaint | Self::Xl | Self::XlInpaint | Self::Turbo => {
+        if use_f16 {
+          "vae/diffusion_pytorch_model.fp16.safetensors"
+        } else {
+          "vae/diffusion_pytorch_model.safetensors"
+        }
+      },
+    }
+  }
+
+  fn clip_file(&self, use_f16: bool) -> &'static str {
+    match self {
+      Self::V1_5 | Self::V1_5Inpaint | Self::V2_1 | Self::V2Inpaint | Self::Xl | Self::XlInpaint | Self::Turbo => {
+        if use_f16 {
+          "text_encoder/model.fp16.safetensors"
+        } else {
+          "text_encoder/model.safetensors"
+        }
+      },
+    }
+  }
+
+  fn clip2_file(&self, use_f16: bool) -> &'static str {
+    match self {
+      Self::V1_5 | Self::V1_5Inpaint | Self::V2_1 | Self::V2Inpaint | Self::Xl | Self::XlInpaint | Self::Turbo => {
+        if use_f16 {
+          "text_encoder_2/model.fp16.safetensors"
+        } else {
+          "text_encoder_2/model.safetensors"
+        }
+      },
+    }
+  }
+}
+
+impl ModelFile {
+  fn get(&self, filename: Option<String>, version: StableDiffusionVersion, use_f16: bool) -> Result<std::path::PathBuf> {
+    use hf_hub::api::sync::Api;
+    match filename {
+      Some(filename) => Ok(std::path::PathBuf::from(filename)),
+      None => {
+        let (repo, path) = match self {
+          Self::Tokenizer => {
+            let tokenizer_repo = match version {
+              StableDiffusionVersion::V1_5 | StableDiffusionVersion::V2_1 | StableDiffusionVersion::V1_5Inpaint | StableDiffusionVersion::V2Inpaint => "openai/clip-vit-base-patch32",
+              StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Turbo => {
+                // This seems similar to the patch32 version except some very small
+                // difference in the split regex.
+                "openai/clip-vit-large-patch14"
+              },
+            };
+            (tokenizer_repo, "tokenizer.json")
+          },
+          Self::Tokenizer2 => ("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", "tokenizer.json"),
+          Self::Clip => (version.repo(), version.clip_file(use_f16)),
+          Self::Clip2 => (version.repo(), version.clip2_file(use_f16)),
+          Self::UnetLcm => ("SimianLuo/LCM_Dreamshaper_v7", if use_f16 { "unet/diffusion_pytorch_model.fp16.safetensors" } else { "unet/diffusion_pytorch_model.safetensors" }),
+          Self::Unet => (version.repo(), version.unet_file(use_f16)),
+          Self::Vae => {
+            // Override for SDXL when using f16 weights.
+            // See https://github.com/huggingface/candle/issues/1060
+            if matches!(version, StableDiffusionVersion::Xl | StableDiffusionVersion::Turbo,) && use_f16 {
+              ("madebyollin/sdxl-vae-fp16-fix", "diffusion_pytorch_model.safetensors")
+            } else {
+              (version.repo(), version.vae_file(use_f16))
+            }
+          },
+        };
+        let filename = Api::new()?.model(repo.to_string()).get(path)?;
+        Ok(filename)
+      },
+    }
+  }
+}
+
+fn output_filename(basename: &str, sample_idx: usize, num_samples: usize, timestep_idx: Option<usize>) -> String {
+  let filename = if num_samples > 1 {
+    match basename.rsplit_once('.') {
+      None => format!("{basename}.{sample_idx}.png"),
+      Some((filename_no_extension, extension)) => {
+        format!("{filename_no_extension}.{sample_idx}.{extension}")
+      },
+    }
+  } else {
+    basename.to_string()
+  };
+  match timestep_idx {
+    None => filename,
+    Some(timestep_idx) => match filename.rsplit_once('.') {
+      None => format!("{filename}-{timestep_idx}.png"),
+      Some((filename_no_extension, extension)) => {
+        format!("{filename_no_extension}-{timestep_idx}.{extension}")
+      },
+    },
+  }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn save_image(vae: &AutoEncoderKL, latents: &Tensor, vae_scale: f64, bsize: usize, idx: usize, final_image: &str, num_samples: usize, timestep_ids: Option<usize>) -> Result<()> {
+  let images = vae.decode(&(latents / vae_scale)?)?;
+  let images = ((images / 2.)? + 0.5)?.to_device(&Device::Cpu)?;
+  let images = (images.clamp(0f32, 1.)? * 255.)?.to_dtype(DType::U8)?;
+  for batch in 0..bsize {
+    let image = images.i(batch)?;
+    let image_filename = output_filename(final_image, (bsize * idx) + batch + 1, batch + num_samples, timestep_ids);
+    candle_examples::save_image(&image, image_filename)?;
+  }
+  Ok(())
+}
+
+#[allow(clippy::too_many_arguments)]
+fn text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option<String>, clip_weights: Option<String>, clip2_weights: Option<String>, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, first: bool) -> Result<Tensor> {
+  let tokenizer_file = if first { ModelFile::Tokenizer } else { ModelFile::Tokenizer2 };
+  let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
+  let tokenizer = Tokenizer::from_file(tokenizer).map_err(E::msg)?;
+  let pad_id = match &sd_config.clip.pad_with {
+    Some(padding) => *tokenizer.get_vocab(true).get(padding.as_str()).unwrap(),
+    None => *tokenizer.get_vocab(true).get("<|endoftext|>").unwrap(),
+  };
+  println!("Running with prompt \"{prompt}\".");
+  let mut tokens = tokenizer.encode(prompt, true).map_err(E::msg)?.get_ids().to_vec();
+  if tokens.len() > sd_config.clip.max_position_embeddings {
+    anyhow::bail!("the prompt is too long, {} > max-tokens ({})", tokens.len(), sd_config.clip.max_position_embeddings)
+  }
+  while tokens.len() < sd_config.clip.max_position_embeddings {
+    tokens.push(pad_id)
+  }
+  let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?;
+
+  println!("Building the Clip transformer.");
+  let clip_weights_file = if first { ModelFile::Clip } else { ModelFile::Clip2 };
+  let clip_weights = if first { clip_weights_file.get(clip_weights, sd_version, use_f16)? } else { clip_weights_file.get(clip2_weights, sd_version, use_f16)? };
+  let clip_config = if first { &sd_config.clip } else { sd_config.clip2.as_ref().unwrap() };
+  let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?;
+  let text_embeddings = text_model.forward(&tokens)?;
+
+  let text_embeddings = if use_guide_scale {
+    let mut uncond_tokens = tokenizer.encode(uncond_prompt, true).map_err(E::msg)?.get_ids().to_vec();
+    if uncond_tokens.len() > sd_config.clip.max_position_embeddings {
+      anyhow::bail!("the negative prompt is too long, {} > max-tokens ({})", uncond_tokens.len(), sd_config.clip.max_position_embeddings)
+    }
+    while uncond_tokens.len() < sd_config.clip.max_position_embeddings {
+      uncond_tokens.push(pad_id)
+    }
+
+    let uncond_tokens = Tensor::new(uncond_tokens.as_slice(), device)?.unsqueeze(0)?;
+    let uncond_embeddings = text_model.forward(&uncond_tokens)?;
+
+    Tensor::cat(&[uncond_embeddings, text_embeddings], 0)?.to_dtype(dtype)?
+  } else {
+    text_embeddings.to_dtype(dtype)?
+  };
+  Ok(text_embeddings)
+}
+
+fn image_preprocess<T: AsRef<std::path::Path>>(path: T) -> anyhow::Result<Tensor> {
+  let img = image::ImageReader::open(path)?.decode()?;
+  let (height, width) = (img.height() as usize, img.width() as usize);
+  let height = height - height % 32;
+  let width = width - width % 32;
+  let img = img.resize_to_fill(width as u32, height as u32, image::imageops::FilterType::CatmullRom);
+  let img = img.to_rgb8();
+  let img = img.into_raw();
+  let img = Tensor::from_vec(img, (height, width, 3), &Device::Cpu)?.permute((2, 0, 1))?.to_dtype(DType::F32)?.affine(2. / 255., -1.)?.unsqueeze(0)?;
+  Ok(img)
+}
+
+/// Convert the mask image to a single channel tensor. Also ensure the image is a multiple of 32 in both dimensions.
+fn mask_preprocess<T: AsRef<std::path::Path>>(path: T) -> anyhow::Result<Tensor> {
+  let img = image::open(path)?.to_luma8();
+  let (new_width, new_height) = {
+    let (width, height) = img.dimensions();
+    (width - width % 32, height - height % 32)
+  };
+  let img = image::imageops::resize(&img, new_width, new_height, image::imageops::FilterType::CatmullRom).into_raw();
+  let mask = Tensor::from_vec(img, (new_height as usize, new_width as usize), &Device::Cpu)?.unsqueeze(0)?.to_dtype(DType::F32)?.div(255.0)?.unsqueeze(0)?;
+  Ok(mask)
+}
+
+/// Generates the mask latents, scaled mask and mask_4 for inpainting. Returns a tuple of None if inpainting is not
+/// being used.
+#[allow(clippy::too_many_arguments)]
+fn inpainting_tensors(sd_version: StableDiffusionVersion, mask_path: Option<String>, dtype: DType, device: &Device, use_guide_scale: bool, vae: &AutoEncoderKL, image: Option<Tensor>, vae_scale: f64) -> Result<(Option<Tensor>, Option<Tensor>, Option<Tensor>)> {
+  match sd_version {
+    StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => {
+      let inpaint_mask = mask_path.ok_or_else(|| anyhow::anyhow!("An inpainting model was requested but mask-path is not provided."))?;
+      // Get the mask image with shape [1, 1, 128, 128]
+      let mask = mask_preprocess(inpaint_mask)?.to_device(device)?.to_dtype(dtype)?;
+      // Generate the masked image from the image and the mask with shape [1, 3, 1024, 1024]
+      let xmask = mask.le(0.5)?.repeat(&[1, 3, 1, 1])?.to_dtype(dtype)?;
+      let image = &image.ok_or_else(|| anyhow::anyhow!("An inpainting model was requested but img2img which is used as the input image is not provided."))?;
+      let masked_img = (image * xmask)?;
+      // Scale down the mask
+      let shape = masked_img.shape();
+      let (w, h) = (shape.dims()[3] / 8, shape.dims()[2] / 8);
+      let mask = mask.interpolate2d(w, h)?;
+      // shape: [1, 4, 128, 128]
+      let mask_latents = vae.encode(&masked_img)?;
+      let mask_latents = (mask_latents.sample()? * vae_scale)?.to_device(device)?;
+
+      let mask_4 = mask.as_ref().repeat(&[1, 4, 1, 1])?;
+      let (mask_latents, mask) = if use_guide_scale { (Tensor::cat(&[&mask_latents, &mask_latents], 0)?, Tensor::cat(&[&mask, &mask], 0)?) } else { (mask_latents, mask) };
+      Ok((Some(mask_latents), Some(mask), Some(mask_4)))
+    },
+    _ => Ok((None, None, None)),
+  }
+}
+
+fn get_tensor_stats(tensor: &Tensor) -> Result<(f32, f32, f32, f32)> {
+  // Convert tensor to F32 for stats calculation - makes statistics more accurate anyway
+  let tensor_f32 = tensor.to_dtype(DType::F32)?;
+  let flat = tensor_f32.flatten_all()?;
+
+  let min = flat.min(0)?.to_scalar::<f32>()?;
+  let max = flat.max(0)?.to_scalar::<f32>()?;
+  let mean = flat.mean_all()?.to_scalar::<f32>()?;
+
+  // Create a tensor with the same shape as flat but filled with the mean value
+  let mean_tensor = Tensor::new(&[mean], flat.device())?;
+  let mean_tensor = mean_tensor.broadcast_as(flat.shape().dims())?;
+  let diff = (flat - mean_tensor)?;
+  let squared = diff.sqr()?;
+  let variance = squared.mean_all()?;
+  let std = variance.sqrt()?.to_scalar::<f32>()?;
+
+  Ok((min, max, mean, std))
+}
+
+// Add this function to track statistics at key points
+fn log_tensor_stats(name: &str, tensor: &Tensor) -> Result<()> {
+  let (min, max, mean, std) = get_tensor_stats(tensor)?;
+  println!("{} stats: min={:.5}, max={:.5}, mean={:.5}, std={:.5}", name, min, max, mean, std);
+  Ok(())
+}
+
+// Add this simpler alternative function
+fn save_tensor_as_numpy(tensor: &Tensor, filename: &str) -> Result<()> {
+  use std::fs::File;
+  use std::io::Write;
+
+  // Move tensor to CPU and convert to f32 for consistent serialization
+  let tensor_cpu = tensor.to_device(&Device::Cpu)?.to_dtype(DType::F32)?;
+  let data = tensor_cpu.to_vec1::<f32>()?;
+
+  // Create a simple binary file with the tensor data
+  let mut file = File::create(filename)?;
+
+  // Write shape information first (as u32 values)
+  let shape = tensor.shape().dims();
+  file.write_all(&(shape.len() as u32).to_le_bytes())?;
+  for &dim in shape {
+    file.write_all(&(dim as u32).to_le_bytes())?;
+  }
+
+  // Write all float values
+  for val in data {
+    file.write_all(&val.to_le_bytes())?;
+  }
+
+  Ok(())
+}
+
+fn run(args: Args) -> Result<()> {
+  use tracing_chrome::ChromeLayerBuilder;
+  use tracing_subscriber::prelude::*;
+
+  let Args { prompt, uncond_prompt, cpu, height, width, n_steps, tokenizer, final_image, sliced_attention_size, num_samples, bsize, sd_version, clip_weights, clip2_weights, vae_weights, unet_weights, tracing, use_f16, guidance_scale, use_flash_attn, mask_path, img2img, strength, seed, only_update_masked, save_tensors, intermediary_images, half_precision } = args;
+
+  if !(0. ..=1.).contains(&strength) {
+    anyhow::bail!("strength should be between 0 and 1, got {strength}")
+  }
+
+  let _guard = if tracing {
+    let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
+    tracing_subscriber::registry().with(chrome_layer).init();
+    Some(guard)
+  } else {
+    None
+  };
+
+  println!("=== LCM-Dreamshaper F16 DEBUG MODE ===");
+  println!("Using dtype: {:?}", if half_precision { DType::F16 } else { DType::F32 });
+
+  let dtype = if half_precision { DType::F16 } else { DType::F32 };
+  let sd_config = match sd_version {
+    StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => stable_diffusion::StableDiffusionConfig::v1_5(sliced_attention_size, height, width),
+    StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint => stable_diffusion::StableDiffusionConfig::v2_1(sliced_attention_size, height, width),
+    StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint => stable_diffusion::StableDiffusionConfig::sdxl(sliced_attention_size, height, width),
+    StableDiffusionVersion::Turbo => stable_diffusion::StableDiffusionConfig::sdxl_turbo(sliced_attention_size, height, width),
+  };
+
+  let mut scheduler = stable_diffusion::lcm::LCMScheduler::new(n_steps, stable_diffusion::lcm::LCMSchedulerConfig::default())?;
+  println!("Scheduler created with dtype={:?}", scheduler.dtype());
+
+  let device = candle_examples::device(cpu)?;
+  // If a seed is not given, generate a random seed and print it
+  let seed = seed.unwrap_or(rand::thread_rng().gen_range(0u64..u64::MAX));
+  println!("Using seed {seed}");
+  device.set_seed(seed)?;
+  let use_guide_scale = guidance_scale > 1.0;
+
+  let which = match sd_version {
+    StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Turbo => vec![true, false],
+    _ => vec![true],
+  };
+
+  // Define the base model and LCM model repositories
+  let base_model_repo = match sd_version {
+    StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => {
+      // Dreamshaper is based on SD 1.5, so use it for those versions
+      "Lykon/dreamshaper-7"
+    },
+    _ => {
+      // For other versions, use the standard repo
+      sd_version.repo()
+    },
+  };
+  let lcm_model_repo = "SimianLuo/LCM_Dreamshaper_v7"; // LCM UNet
+
+  // Create a custom ModelFile enum variant for the base model
+  let clip_weights = if clip_weights.is_none() {
+    use hf_hub::api::sync::Api;
+    let api = Api::new()?;
+    let clip_path = api.model(base_model_repo.to_string()).get(if use_f16 { "text_encoder/model.fp16.safetensors" } else { "text_encoder/model.safetensors" })?;
+    Some(clip_path.to_string_lossy().to_string())
+  } else {
+    clip_weights
+  };
+
+  let vae_weights = if vae_weights.is_none() {
+    use hf_hub::api::sync::Api;
+    let api = Api::new()?;
+    let vae_path = api.model(base_model_repo.to_string()).get(if use_f16 { "vae/diffusion_pytorch_model.fp16.safetensors" } else { "vae/diffusion_pytorch_model.safetensors" })?;
+    Some(vae_path.to_string_lossy().to_string())
+  } else {
+    vae_weights
+  };
+
+  // For the UNet, always use the LCM model
+  let unet_weights = {
+    use hf_hub::api::sync::Api;
+    let api = Api::new()?;
+    let unet_path = api.model(lcm_model_repo.to_string()).get(if use_f16 {
+      if lcm_model_repo == "SimianLuo/LCM_Dreamshaper_v7" {
+        "unet/diffusion_pytorch_model.safetensors"
+      } else {
+        "unet/diffusion_pytorch_model.fp16.safetensors"
+      }
+    } else {
+      "unet/diffusion_pytorch_model.safetensors"
+    })?;
+    Some(unet_path.to_string_lossy().to_string())
+  };
+
+  // Load text embeddings from the base model
+  let text_embeddings = which.iter().map(|first| text_embeddings(&prompt, &uncond_prompt, tokenizer.clone(), clip_weights.clone(), clip2_weights.clone(), sd_version, &sd_config, use_f16, &device, dtype, use_guide_scale, *first)).collect::<Result<Vec<_>>>()?;
+
+  let text_embeddings = Tensor::cat(&text_embeddings, D::Minus1)?;
+  let text_embeddings = text_embeddings.repeat((bsize, 1, 1))?;
+  println!("Text embeddings shape: {:?}", text_embeddings.shape());
+  if PRINT_DEBUG_TENSORS {
+    println!("Text embeddings first few values: {}", text_embeddings.i(0)?.i(0)?.narrow(0, 0, 10)?);
+  }
+
+  // Replace the random projection with a fixed one to match diffusers
+  let text_embeddings = if text_embeddings.dim(2)? == 2048 {
+    println!("Projecting text embeddings from 2048 to 768 dimensions");
+
+    // Load the projection matrix from a file (you'll need to extract this from the diffusers model)
+    // For now, we'll still use a random one but with a fixed seed
+    device.set_seed(42)?; // Use a fixed seed for the projection
+    let projection = Tensor::randn(0.0, 0.02, (2048, 768), &device)?.to_dtype(dtype)?;
+
+    let proj_embeddings = text_embeddings.matmul(&projection)?;
+    proj_embeddings
+  } else {
+    text_embeddings
+  };
+  if PRINT_DEBUG_TENSORS {
+    println!("Projected text embeddings: {:?}", text_embeddings);
+  }
+
+  println!("Building the autoencoder.");
+  let vae_weights = ModelFile::Vae.get(vae_weights, sd_version, use_f16)?;
+  let vae = sd_config.build_vae(vae_weights, &device, dtype)?;
+
+  let (image, init_latent_dist) = match &img2img {
+    None => (None, None),
+    Some(image) => {
+      let image = image_preprocess(image)?.to_device(&device)?.to_dtype(dtype)?;
+      println!("Input image shape: {:?}", image.shape());
+      (Some(image.clone()), Some(vae.encode(&image)?))
+    },
+  };
+
+  println!(
+    "Latent shape after VAE encoding: {:?}",
+    init_latent_dist.as_ref().map(|dist| {
+      // Try to get shape from a temporary sample
+      match dist.sample() {
+        Ok(sample) => format!("{:?}", sample.shape()),
+        Err(_) => "Unable to determine shape".to_string(),
+      }
+    })
+  );
+
+  println!("Building the unet.");
+  let unet_weights = ModelFile::UnetLcm.get(unet_weights, sd_version, use_f16)?;
+  let in_channels = match sd_version {
+    StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => 9,
+    _ => 4,
+  };
+  println!("Loading UNet weights from: SimianLuo/LCM_Dreamshaper_v7/unet");
+
+  // Ensure we're loading the correct UNet configuration
+  println!("Building UNet for LCM (SimianLuo/LCM_Dreamshaper_v7)");
+
+  // For LCM models, ensure that:
+  // 1. We're using the correct prediction type (epsilon)
+  // 2. The cross-attention dimension is correct (768 for Dreamshaper)
+  // 3. The in_channels matches the expected input (4 for standard, 9 for inpainting)
+  let unet_config = UNet2DConditionModelConfig {
+    // LCM specific config values
+    cross_attention_dim: 768, // Dreamshaper v7 uses 768
+    // Rest of configuration...
+    center_input_sample: false,
+    flip_sin_to_cos: true,
+    freq_shift: 0.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.0,
+    norm_num_groups: 32,
+    norm_eps: 1e-5,
+    sliced_attention_size: None,
+    use_linear_projection: false,
+  };
+
+  // Load the model directly from the safetensors file
+  let vs_unet = unsafe { VarBuilder::from_mmaped_safetensors(&[unet_weights], dtype, &device)? };
+
+  let unet = UNet2DConditionModel::new(vs_unet, in_channels, 4, use_flash_attn, unet_config)?;
+
+  // Ensure strength is at least 0.4 for img2img to work well
+  let effective_strength = if img2img.is_some() {
+    strength.max(0.4).min(0.99) // Keep between 0.4 and 0.99
+  } else {
+    strength
+  };
+
+  // Make sure we're doing a proper image-to-image transform
+  // with the right amount of noise
+  println!("Using img2img with strength={:.4}", effective_strength);
+  let mut raw_t_start = n_steps - (n_steps as f64 * effective_strength).ceil() as usize;
+
+  let vae_scale = match sd_version {
+    StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint | StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Xl => 0.18215,
+    StableDiffusionVersion::Turbo => 0.13025,
+  };
+
+  let (mask_latents, mask, mask_4) = inpainting_tensors(sd_version, mask_path, dtype, &device, use_guide_scale, &vae, image, vae_scale)?;
+
+  println!("Timesteps: {:?}", scheduler.timesteps());
+  // println!("Alpha cumprod: {:?}", scheduler.alphas_cumprod());
+
+  // Enable debug dumps to compare with diffusers
+  let enable_debug_dumps = true;
+  let debug_dir = "lcm_rust_debug";
+  if enable_debug_dumps {
+    std::fs::create_dir_all(debug_dir).expect("Failed to create debug directory");
+
+    // Log the scheduler configuration
+    let mut config_file = std::fs::File::create(format!("{}/scheduler_config.txt", debug_dir)).expect("Failed to create config file");
+    use std::io::Write;
+    writeln!(config_file, "Rust LCM Configuration:").unwrap();
+    writeln!(config_file, "  beta_start: {}", scheduler.config.beta_start).unwrap();
+    writeln!(config_file, "  beta_end: {}", scheduler.config.beta_end).unwrap();
+    // More config params...
+    writeln!(config_file, "\nTimesteps: {:?}", scheduler.timesteps()).unwrap();
+    writeln!(config_file, "First 10 alphas_cumprod: {:?}", scheduler.alphas_cumprod().iter().take(10).collect::<Vec<_>>()).unwrap();
+  }
+
+  // Add this code before the loop where we iterate through samples
+  // This will generate and log the guidance scale embedding
+
+  // Create and log guidance scale embedding for comparison with Python
+  println!("\n==== Generating LCM Guidance Scale Embedding ====");
+  let guidance_scale = args.guidance_scale;
+  println!("Guidance scale: {}", guidance_scale);
+
+  // The embedding dimension should match the UNet time embedding dimension
+  // For LCM Dreamshaper, this is typically 1280 (not 320)
+  // Use a hardcoded value since we don't have direct access to the config
+  let embedding_dim = 1280; // Standard for SD models based on UNet config
+  println!("Embedding dimension: {}", embedding_dim);
+
+  // Generate the embedding
+  let guidance_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &device, dtype)?;
+
+  println!("Guidance embedding shape: {:?}", guidance_embedding.shape());
+
+  // Save the embedding to a numpy file for comparison with Python
+  if save_tensors {
+    let tensor_dir = "lcm_rust_tensors";
+    std::fs::create_dir_all(tensor_dir)?;
+
+    save_tensor_as_numpy(&guidance_embedding, &format!("{}/guidance_embedding.npy", tensor_dir))?;
+    println!("Saved guidance embedding to {}/guidance_embedding.npy", tensor_dir);
+  }
+  println!("===============================================\n");
+
+  for idx in 0..num_samples {
+    let timesteps = scheduler.timesteps().to_vec();
+
+    // Initialize latents
+    let latents = if let Some(init_latent_dist) = &init_latent_dist {
+      println!("Using img2img with strength={:.4}", effective_strength);
+
+      // LCM can be more sensitive to the starting timestep
+      // Get the latents from the VAE
+      let latents = (init_latent_dist.sample()? * vae_scale)?.to_device(&device)?;
+
+      // Calculate the actual start timestep based on strength
+      // In LCM, starting too late can cause artifacts, so ensure we start early enough
+      let t_start_index = (timesteps.len() as f64 * (1.0 - effective_strength.max(0.5))).floor() as usize;
+      let t_start = timesteps[t_start_index.min(timesteps.len() - 1)];
+
+      println!("Starting from timestep {}/{} (index: {})", t_start, timesteps[0], t_start_index);
+
+      // Use a fixed seed for noise to improve reproducibility
+      device.set_seed(seed)?;
+      let noise = (Tensor::randn(0f32, 1f32, latents.shape(), &device)?).to_dtype(dtype)?;
+
+      // Clone the noise when passing it to add_noise
+      let noised_latents = scheduler.add_noise(&latents, noise.clone(), t_start)?;
+
+      // Now we can use noise for logging since we cloned it above
+      log_tensor_stats("initial_latents", &latents)?;
+      log_tensor_stats("noise", &noise)?;
+      log_tensor_stats("noised_latents", &noised_latents)?;
+
+      // Check the dtype of noised_latents
+      debug_tensor("final_img2img_latents", &noised_latents)?;
+
+      // Explicitly convert to expected dtype to avoid errors
+      noised_latents.to_dtype(dtype)?
+    } else {
+      // Regular text-to-image generation with random noise
+      let latents = Tensor::randn(0f32, 1f32, (bsize, 4, sd_config.height / 8, sd_config.width / 8), &device)?;
+
+      // Use affine for multiplication
+      let scaled_latents = latents.affine(scheduler.init_noise_sigma(), 0.0)?;
+      debug_tensor("scaled_initial_latents", &scaled_latents)?;
+
+      // Explicitly convert to expected dtype
+      scaled_latents.to_dtype(dtype)?
+    };
+    let mut latents = latents.to_dtype(dtype)?;
+
+    println!("Setting seed to {}", seed);
+    device.set_seed(seed)?;
+
+    if PRINT_DEBUG_TENSORS {
+      println!("Initial latents shape: {:?}", latents.shape());
+      println!("Initial latents first few values: {}", latents.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
+    }
+
+    println!("Starting LCM sampling");
+    let start_time = std::time::Instant::now();
+
+    for (timestep_index, &timestep) in timesteps.iter().enumerate() {
+      if timestep_index < raw_t_start {
+        continue;
+      }
+
+      println!("\n--- Timestep {}/{} (value: {}) ---", timestep_index + 1, timesteps.len(), timestep);
+
+      let latent_model_input = if use_guide_scale {
+        println!("Using classifier-free guidance with scale: {}", guidance_scale);
+        Tensor::cat(&[&latents, &latents], 0)?
+      } else {
+        latents.clone()
+      };
+
+      if PRINT_DEBUG_TENSORS {
+        println!("Latent input shape before scaling: {:?}", latent_model_input.shape());
+        println!("Latent input first few values: {}", latent_model_input.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
+      }
+
+      let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep);
+
+      if PRINT_DEBUG_TENSORS {
+        println!("Latent input shape after scaling: {:?}", latent_model_input.shape());
+        println!("Scaled latent input first few values: {}", latent_model_input.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
+      }
+
+      // Generate guidance scale embedding
+      let guidance_emb = if use_guide_scale {
+        // For LCM/Dreamshaper, the time embedding dimension is 1280
+        let time_embed_dim = 1280; // Standard for SD models (320 base channels * 4)
+
+        println!("Generating guidance scale embedding (scale={})", guidance_scale);
+        Some(scheduler.get_guidance_scale_embedding(guidance_scale, time_embed_dim, &device, dtype)?)
+      } else {
+        None
+      };
+
+      // Before UNet forward call
+      debug_tensor("latent_input_to_unet", &latent_model_input)?;
+      debug_tensor("text_embeddings_to_unet", &text_embeddings)?;
+
+      // Pass the guidance embedding to the UNet
+      // This would require modifying the UNet forward method to accept this parameter
+      println!("Running UNet forward pass with guidance embedding...");
+      let noise_pred = unet.forward_with_guidance(&latent_model_input, timestep as f64, &text_embeddings, guidance_emb.as_ref())?;
+
+      if PRINT_DEBUG_TENSORS {
+        println!("Noise prediction shape: {:?}", noise_pred.shape());
+        println!("Noise prediction first few values: {}", noise_pred.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
+      }
+
+      let noise_pred = if use_guide_scale {
+        let noise_pred = noise_pred.chunk(2, 0)?;
+        let (noise_pred_uncond, noise_pred_text) = (&noise_pred[0], &noise_pred[1]);
+
+        // LCM requires a different guidance approach than standard diffusion
+        // Linear interpolation with scale factor
+        let diff = (noise_pred_text - noise_pred_uncond)?;
+        let scaled_diff = diff.affine(guidance_scale, 0.0)?;
+        (noise_pred_uncond + &scaled_diff)?
+      } else {
+        noise_pred
+      };
+
+      if PRINT_DEBUG_TENSORS {
+        println!("Latents before step first few values: {}", latents.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
+      }
+
+      println!("Applying scheduler step...");
+      latents = scheduler.step(&noise_pred, timestep, &latents)?;
+
+      if PRINT_DEBUG_TENSORS {
+        println!("Latents after step first few values: {}", latents.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
+      }
+
+      println!("Step {}/{} completed", timestep_index + 1, n_steps);
+
+      // Then use it at key points
+      log_tensor_stats("latents_before_step", &latents)?;
+      log_tensor_stats("noise_pred", &noise_pred)?;
+      log_tensor_stats("latents_after_step", &latents)?;
+
+      if save_tensors {
+        let tensor_dir = "lcm_rust_tensors";
+        std::fs::create_dir_all(tensor_dir)?;
+
+        // Save key tensors at each step
+        let step_dir = format!("{}/step_{}", tensor_dir, timestep_index);
+        std::fs::create_dir_all(&step_dir)?;
+
+        // Save various tensors as numpy arrays for later comparison
+        save_tensor_as_numpy(&latent_model_input, &format!("{}/input.npy", step_dir))?;
+        save_tensor_as_numpy(&noise_pred, &format!("{}/noise_pred.npy", step_dir))?;
+        save_tensor_as_numpy(&latents, &format!("{}/output.npy", step_dir))?;
+      }
+
+      // After UNet forward call
+      debug_tensor("noise_pred_from_unet", &noise_pred)?;
+    }
+
+    let dt = start_time.elapsed().as_secs_f32();
+    println!("Sampling completed in {:.2}s", dt);
+
+    // Generate final image
+    println!("Generating final image for sample {}/{}", idx + 1, num_samples);
+    save_image(&vae, &latents, vae_scale, bsize, idx, &final_image, num_samples, None)?;
+  }
+  Ok(())
+}
+
+// Add this debug function to check dtype at critical points
+fn debug_tensor(name: &str, tensor: &Tensor) -> Result<()> {
+  println!("DEBUG {}: shape={:?}, dtype={:?}", name, tensor.shape(), tensor.dtype());
+  Ok(())
+}
+
+fn main() -> Result<()> {
+  let args = Args::parse();
+  run(args)
+}
diff --git a/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/lcm.rs b/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/lcm.rs
new file mode 100644
index 0000000000..8c0b7d27ef
--- /dev/null
+++ b/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/lcm.rs
@@ -0,0 +1,762 @@
+use super::schedulers::{betas_for_alpha_bar, BetaSchedule, PredictionType, TimestepSpacing};
+use candle::{Result, Tensor, IndexOp, Device, DType, Shape};
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum LCMVarianceType {
+  FixedSmall,
+  FixedSmallLog,
+  FixedLarge,
+  FixedLargeLog,
+  Learned,
+}
+
+impl Default for LCMVarianceType {
+  fn default() -> Self {
+    Self::FixedSmall
+  }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum LCMSolverType {
+  Default,
+  LogRho,
+}
+
+impl Default for LCMSolverType {
+  fn default() -> Self {
+    Self::Default
+  }
+}
+
+#[derive(Debug, Clone)]
+pub struct LCMSchedulerConfig {
+  /// The value of beta at the beginning of training.
+  pub beta_start: f64,
+  /// The value of beta at the end of training.
+  pub beta_end: f64,
+  /// How beta evolved during training.
+  pub beta_schedule: BetaSchedule,
+  /// Option to predicted sample between -1 and 1 for numerical stability.
+  pub clip_sample: bool,
+  /// Option to clip the variance used when adding noise to the denoised sample.
+  pub variance_type: LCMVarianceType,
+  /// prediction type of the scheduler function
+  pub prediction_type: PredictionType,
+  /// number of diffusion steps used to train the model.
+  pub train_timesteps: usize,
+  /// time step spacing for the diffusion process
+  pub timestep_spacing: TimestepSpacing,
+  /// Original training steps
+  pub original_inference_steps: usize,
+  /// Timestep scaling factor (default: 10.0 in diffusers)
+  pub timestep_scaling: f64,
+  /// Adjust the indexes of the inference schedule by this value.
+  pub steps_offset: usize,
+  /// The solver type to use for the diffusion process
+  pub solver_type: LCMSolverType,
+}
+
+impl Default for LCMSchedulerConfig {
+  fn default() -> Self {
+    Self { beta_start: 0.00085, beta_end: 0.012, beta_schedule: BetaSchedule::ScaledLinear, clip_sample: false, variance_type: LCMVarianceType::FixedSmall, prediction_type: PredictionType::Epsilon, train_timesteps: 1000, timestep_spacing: TimestepSpacing::Leading, original_inference_steps: 50, timestep_scaling: 10.0, steps_offset: 1, solver_type: LCMSolverType::Default }
+  }
+}
+
+pub struct LCMScheduler {
+  alphas_cumprod: Vec<f64>,
+  init_noise_sigma: f64,
+  timesteps: Vec<usize>,
+  step_ratio: usize,
+  pub config: LCMSchedulerConfig,
+  step_index: Option<usize>,
+  begin_index: Option<usize>,
+  dtype: DType,
+}
+
+impl LCMScheduler {
+  pub fn new(inference_steps: usize, config: LCMSchedulerConfig) -> Result<Self> {
+    // Set up the correct betas
+    let betas = match config.beta_schedule {
+      BetaSchedule::ScaledLinear => {
+        // LCM uses scaled_linear specifically this way
+        let start = config.beta_start.sqrt();
+        let end = config.beta_end.sqrt();
+        let betas = super::utils::linspace(start, end, config.train_timesteps)?;
+        betas.sqr()?
+      },
+      // Other schedules...
+      _ => unreachable!("LCM needs ScaledLinear"),
+    };
+
+    // Set up alphas exactly as in diffusers
+    let betas = betas.to_vec1::<f64>()?;
+    let alphas = betas.iter().map(|beta| 1.0 - beta).collect::<Vec<_>>();
+
+    let mut alphas_cumprod = Vec::with_capacity(alphas.len());
+    let mut alpha_prod = 1.0;
+    for alpha in alphas.iter() {
+      alpha_prod *= alpha;
+      alphas_cumprod.push(alpha_prod);
+    }
+
+    // Create timesteps based on TimestepSpacing, similar to DDIM
+    let step_ratio = config.train_timesteps / inference_steps;
+    let timesteps: Vec<usize> = match config.timestep_spacing {
+      TimestepSpacing::Leading => (0..(inference_steps)).map(|s| s * step_ratio + config.steps_offset).rev().collect(),
+      TimestepSpacing::Trailing => std::iter::successors(Some(config.train_timesteps), |n| if *n > step_ratio { Some(n - step_ratio) } else { None }).map(|n| n - 1).collect(),
+      TimestepSpacing::Linspace => super::utils::linspace(0.0, (config.train_timesteps - 1) as f64, inference_steps)?.to_vec1::<f64>()?.iter().map(|&f| f as usize).rev().collect(),
+    };
+
+    println!("Using LCM timesteps with {:?} spacing: {:?}", config.timestep_spacing, timesteps);
+
+    let init_noise_sigma = 1.0;
+
+    Ok(Self { alphas_cumprod, init_noise_sigma, timesteps, step_ratio, config, step_index: None, begin_index: None, dtype: DType::F32 })
+  }
+
+  pub fn with_dtype(mut self, dtype: DType) -> Self {
+    self.dtype = dtype;
+    self
+  }
+
+  fn get_variance(&self, timestep: usize) -> f64 {
+    let prev_t = timestep as isize - self.step_ratio as isize;
+    let alpha_prod_t = self.alphas_cumprod[timestep];
+    let alpha_prod_t_prev = if prev_t >= 0 { self.alphas_cumprod[prev_t as usize] } else { 1.0 };
+    let current_beta_t = 1. - alpha_prod_t / alpha_prod_t_prev;
+
+    // For t > 0, compute predicted variance βt (see formula (6) and (7) from [the pdf](https://arxiv.org/pdf/2006.11239.pdf))
+    // and sample from it to get previous sample
+    // x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
+    let variance = (1. - alpha_prod_t_prev) / (1. - alpha_prod_t) * current_beta_t;
+
+    // retrieve variance
+    match self.config.variance_type {
+      LCMVarianceType::FixedSmall => variance.max(1e-20),
+      // for rl-diffuser https://arxiv.org/abs/2205.09991
+      LCMVarianceType::FixedSmallLog => {
+        let variance = variance.max(1e-20).ln();
+        (variance * 0.5).exp()
+      },
+      LCMVarianceType::FixedLarge => current_beta_t,
+      LCMVarianceType::FixedLargeLog => current_beta_t.ln(),
+      LCMVarianceType::Learned => variance,
+    }
+  }
+
+  pub fn timesteps(&self) -> &[usize] {
+    self.timesteps.as_slice()
+  }
+
+  ///  Ensures interchangeability with schedulers that need to scale the denoising model input
+  /// depending on the current timestep.
+  pub fn scale_model_input(&self, sample: Tensor, timestep: usize) -> Tensor {
+    // LCM doesn't apply any special scaling to model inputs
+    println!("LCM: No scaling applied to model input for timestep {}", timestep);
+    sample
+  }
+
+  fn get_scalings_for_boundary_condition_discrete(&self, timestep: usize) -> (f64, f64) {
+    // LCM uses a fixed sigma_data value of 0.5
+    let sigma_data = 0.5f64;
+
+    // Timestep scaling factor from diffusers configuration
+    let timestep_scaling = self.config.timestep_scaling;
+
+    // In diffusers, they don't use the timestep value directly
+    let alpha_cumprod = self.alphas_cumprod[timestep];
+    let sqrt_alpha_prod = alpha_cumprod.sqrt();
+    let sqrt_one_minus_alpha_prod = (1.0 - alpha_cumprod).sqrt();
+
+    // This is the key calculation - convert alpha to sigma space with scaling
+    let scaled_timestep = sqrt_one_minus_alpha_prod / sqrt_alpha_prod * timestep_scaling;
+
+    println!("Boundary condition for timestep={}: scaled_t={} (with scaling={})", timestep, scaled_timestep, timestep_scaling);
+
+    // Calculate boundary condition coefficients with scaled timestep
+    let c_skip = sigma_data.powf(2.0) / (scaled_timestep.powf(2.0) + sigma_data.powf(2.0));
+    let c_out = scaled_timestep / (scaled_timestep.powf(2.0) + sigma_data.powf(2.0)).sqrt();
+
+    (c_skip, c_out)
+  }
+
+  pub fn step(&mut self, model_output: &Tensor, timestep: usize, sample: &Tensor) -> Result<Tensor> {
+    // Initialize step_index if it's not already set
+    if self.step_index.is_none() {
+      self.init_step_index(timestep)?;
+      println!("Initialized step_index to {} for timestep {}", self.step_index.unwrap(), timestep);
+    }
+
+    // Get current timestep index
+    let timestep_index = self.step_index.unwrap();
+
+    // Calculate prev_timestep_index
+    let prev_timestep_index = timestep_index + 1;
+    let prev_timestep = if prev_timestep_index >= self.timesteps.len() {
+      0 // Final step
+    } else {
+      self.timesteps[prev_timestep_index]
+    };
+
+    // Update step_index for next time
+    self.step_index = Some(prev_timestep_index);
+
+    // Choose solver based on configuration
+    match self.config.solver_type {
+      LCMSolverType::Default => self.step_default(model_output, timestep, prev_timestep, sample),
+      LCMSolverType::LogRho => self.step_logrho(model_output, timestep, prev_timestep, sample),
+    }
+  }
+
+  // Default solver (existing logic moved to a separate method)
+  fn step_default(&self, model_output: &Tensor, timestep: usize, prev_timestep: usize, sample: &Tensor) -> Result<Tensor> {
+    // Get dtype from sample and use it for consistency
+    let target_dtype = sample.dtype();
+    let device = sample.device();
+
+    println!("Step (Default): current_t={}, prev_t={}", timestep, prev_timestep);
+
+    // Convert model output to match sample dtype
+    let model_output = model_output.to_dtype(target_dtype)?;
+
+    // Get the exact alpha values for this timestep
+    let alpha_prod_t = if timestep < self.alphas_cumprod.len() {
+      self.alphas_cumprod[timestep]
+    } else {
+      println!("Warning: Timestep {} out of range, using last alpha", timestep);
+      *self.alphas_cumprod.last().unwrap()
+    };
+
+    // Get the previous alpha value
+    let alpha_prod_t_prev = if prev_timestep == 0 {
+      1.0 // Final step goes to alpha=1
+    } else if prev_timestep < self.alphas_cumprod.len() {
+      self.alphas_cumprod[prev_timestep]
+    } else {
+      println!("Warning: Prev timestep {} out of range", prev_timestep);
+      *self.alphas_cumprod.last().unwrap()
+    };
+
+    // Calculate coefficients for boundary conditions
+    let (c_skip, c_out) = self.get_scalings_for_boundary_condition_discrete(timestep);
+
+    // Compute predicted original sample from noise prediction
+    let pred_original_sample = match self.config.prediction_type {
+      PredictionType::Epsilon => {
+        let beta_prod_t = 1.0 - alpha_prod_t;
+        let beta_prod_t_sqrt = beta_prod_t.sqrt();
+
+        let scaled_noise = model_output.affine(beta_prod_t_sqrt, 0.0)?;
+        let diff = (sample - &scaled_noise)?;
+        diff.affine(1.0 / alpha_prod_t.sqrt(), 0.0)?
+      },
+      PredictionType::Sample => model_output.clone(),
+      PredictionType::VPrediction => {
+        let sqrt_alpha_prod = alpha_prod_t.sqrt();
+        let sqrt_one_minus_alpha_prod = (1.0 - alpha_prod_t).sqrt();
+
+        let term1 = sample.affine(sqrt_alpha_prod, 0.0)?;
+        let term2 = model_output.affine(sqrt_one_minus_alpha_prod, 0.0)?;
+        (term1 - term2)?
+      },
+      _ => unreachable!(),
+    };
+
+    // Apply possible clipping
+    let pred_orig = if self.config.clip_sample { pred_original_sample.clamp(-1f32, 1f32)? } else { pred_original_sample };
+
+    // Apply boundary condition
+    let term1 = pred_orig.affine(c_out, 0.0)?;
+    let term2 = sample.affine(c_skip, 0.0)?;
+    let denoised = (&term1 + &term2)?;
+
+    // Set up the previous sample
+    let pred_prev_sample = if prev_timestep == 0 {
+      denoised.clone()
+    } else {
+      let prev_sqrt_alpha = alpha_prod_t_prev.sqrt();
+      denoised.affine(prev_sqrt_alpha, 0.0)?
+    };
+
+    Ok(pred_prev_sample)
+  }
+
+  // LogRho solver implementation
+  fn step_logrho(&self, model_output: &Tensor, timestep: usize, prev_timestep: usize, sample: &Tensor) -> Result<Tensor> {
+    // Ensure consistent dtype by getting it from the sample
+    let target_dtype = sample.dtype();
+    let device = sample.device();
+
+    println!("Step (LogRho): current_t={}, prev_t={}", timestep, prev_timestep);
+    println!("  Using target dtype: {:?}", target_dtype);
+
+    // Convert model output to match sample dtype
+    let model_output = model_output.to_dtype(target_dtype)?;
+    println!("  model_output dtype after conversion: {:?}", model_output.dtype());
+
+    // Get alpha values
+    let alpha_prod_t = if timestep < self.alphas_cumprod.len() { self.alphas_cumprod[timestep] } else { *self.alphas_cumprod.last().unwrap() };
+
+    let alpha_prod_t_prev = if prev_timestep == 0 {
+      1.0
+    } else if prev_timestep < self.alphas_cumprod.len() {
+      self.alphas_cumprod[prev_timestep]
+    } else {
+      *self.alphas_cumprod.last().unwrap()
+    };
+
+    // Special case: if we're at the last timestep, return the predicted sample directly
+    if prev_timestep == 0 {
+      // Handle final step differently - exact same implementation as in default solver
+      return self.step_default(&model_output, timestep, prev_timestep, sample);
+    }
+
+    // Calculate sigma values for LogRho solver with numerical stability
+    let sigma_t = ((1.0 - alpha_prod_t).max(1e-12)).sqrt();
+    let sigma_t_prev = ((1.0 - alpha_prod_t_prev).max(1e-12)).sqrt();
+
+    // Calculate key logrho parameters with numerical stability
+    let logrho_t = sigma_t.ln() - alpha_prod_t.sqrt().max(1e-12).ln();
+    let logrho_t_prev = sigma_t_prev.ln() - alpha_prod_t_prev.sqrt().max(1e-12).ln();
+
+    // Calculate h parameter (step size in log rho space)
+    let h = logrho_t_prev - logrho_t;
+    let exp_h = h.exp().clamp(0.01, 5.0); // Constrain to a reasonable range
+    println!("  LogRho params: h={:.6}, exp(h)={:.6}", h, exp_h);
+
+    // For epsilon prediction we need to transform the model output
+    let score = match self.config.prediction_type {
+      PredictionType::Epsilon => {
+        // For ε prediction type, score is -σ/√α_t * ε
+        let scale = -sigma_t / alpha_prod_t.sqrt().max(1e-12);
+        model_output.affine(scale, 0.0)?.to_dtype(target_dtype)?
+      },
+      PredictionType::VPrediction => {
+        // model_output is already in the right form
+        model_output.clone()
+      },
+      PredictionType::Sample => {
+        // Score = (x_t - √α_t * x_0) / σ_t
+        let alpha_sqrt = alpha_prod_t.sqrt();
+        let scaled_pred = model_output.affine(alpha_sqrt, 0.0)?.to_dtype(target_dtype)?;
+        let diff = (sample.clone() - scaled_pred)?.to_dtype(target_dtype)?;
+        diff.affine(1.0 / sigma_t, 0.0)?.to_dtype(target_dtype)?
+      },
+      _ => unreachable!(),
+    };
+
+    println!("  score dtype: {:?}", score.dtype());
+
+    // Calculate d term: d = x_t + σ_t^2/√α_t * score
+    let scale_d = sigma_t.powi(2) / alpha_prod_t.sqrt().max(1e-12);
+    let score_scaled = score.affine(scale_d, 0.0)?.to_dtype(target_dtype)?;
+    let d = (sample + score_scaled)?.to_dtype(target_dtype)?;
+    println!("  d dtype: {:?}", d.dtype());
+
+    // Apply the main update formula directly:
+    // x_prev = exp_h * sample + sigma_prev * (1 - exp_h) * d
+
+    // First part: exp_h * sample
+    let term1 = sample.affine(exp_h, 0.0)?.to_dtype(target_dtype)?;
+    println!("  term1 dtype: {:?}", term1.dtype());
+
+    // Second part: sigma_prev * (1 - exp_h) * d
+    let term2_factor = sigma_t_prev * (1.0 - exp_h);
+    println!("  term2_factor: {:.6}", term2_factor);
+
+    let term2 = d.affine(term2_factor, 0.0)?.to_dtype(target_dtype)?;
+    println!("  term2 dtype: {:?}", term2.dtype());
+
+    // Add the terms together
+    let pred_prev_sample = (&term1 + &term2)?.to_dtype(target_dtype)?;
+    println!("  pred_prev_sample dtype: {:?}", pred_prev_sample.dtype());
+
+    Ok(pred_prev_sample)
+  }
+
+  pub fn add_noise(&self, original_samples: &Tensor, noise: Tensor, timestep: usize) -> Result<Tensor> {
+    // Get the correct alpha value
+    let alpha_cumprod = self.alphas_cumprod[timestep];
+
+    println!("Adding noise for timestep={}, alpha={:.6}, sqrt(alpha)={:.6}", timestep, alpha_cumprod, alpha_cumprod.sqrt());
+    println!("  original_samples: shape={:?}, dtype={:?}", original_samples.shape(), original_samples.dtype());
+    println!("  noise: shape={:?}, dtype={:?}", noise.shape(), noise.dtype());
+
+    // Keep operations in original_samples' dtype for now
+    let target_dtype = original_samples.dtype();
+    let noise = noise.to_dtype(target_dtype)?;
+
+    // Calculate coefficients
+    let sqrt_alpha = alpha_cumprod.sqrt();
+    let sqrt_one_minus_alpha = (1.0 - alpha_cumprod).sqrt();
+
+    // Apply noise formula
+    let scaled_samples = original_samples.affine(sqrt_alpha, 0.0)?;
+    let scaled_noise = noise.affine(sqrt_one_minus_alpha, 0.0)?;
+
+    println!("  scaled_samples: shape={:?}, dtype={:?}", scaled_samples.shape(), scaled_samples.dtype());
+    println!("  scaled_noise: shape={:?}, dtype={:?}", scaled_noise.shape(), scaled_noise.dtype());
+
+    // Add scaled noise - should maintain dtype
+    let result = scaled_samples.add(&scaled_noise)?;
+
+    println!("  result: shape={:?}, dtype={:?}", result.shape(), result.dtype());
+
+    Ok(result)
+  }
+
+  pub fn init_noise_sigma(&self) -> f64 {
+    self.init_noise_sigma
+  }
+
+  pub fn alphas_cumprod(&self) -> &[f64] {
+    &self.alphas_cumprod
+  }
+
+  pub fn log_config(&self) {
+    println!("LCM Scheduler Configuration:");
+    println!("  beta_start: {}", self.config.beta_start);
+    println!("  beta_end: {}", self.config.beta_end);
+    println!("  beta_schedule: {:?}", self.config.beta_schedule);
+    println!("  train_timesteps: {}", self.config.train_timesteps);
+    println!("  prediction_type: {:?}", self.config.prediction_type);
+    println!("  clip_sample: {}", self.config.clip_sample);
+    println!("  variance_type: {:?}", self.config.variance_type);
+    println!("  timesteps: {:?}", self.timesteps);
+    // Print the first few alphas_cumprod values
+    let num_to_show = std::cmp::min(10, self.alphas_cumprod.len());
+    println!("  First {} alphas_cumprod: {:?}", num_to_show, self.alphas_cumprod.iter().take(num_to_show).collect::<Vec<_>>());
+  }
+
+  /// Creates embeddings from guidance scale values, similar to timestep embeddings
+  /// Based on diffusers implementation
+  pub fn get_guidance_scale_embedding(&self, guidance_scale: f64, embedding_dim: usize, device: &Device, dtype: DType) -> Result<Tensor> {
+    // Scale by exactly 1000.0 to match diffusers
+    let w = guidance_scale * 1000.0;
+    println!("Creating guidance embedding with dim={}, scaled value={}", embedding_dim, w);
+
+    // Create embedding vector directly to avoid tensor operations
+    let mut embedding = Vec::with_capacity(embedding_dim);
+    let half_dim = embedding_dim / 2;
+
+    // Compute the position frequencies
+    for i in 0..half_dim {
+      let freq = (i as f32) * (-((10000.0f32).ln()) / ((half_dim - 1) as f32));
+      let freq = freq.exp();
+
+      // Apply the guidance scale (w value)
+      let value = w as f32 * freq;
+
+      // Add sin and cos values
+      embedding.push(value.sin());
+      embedding.push(value.cos());
+    }
+
+    // If odd dimension, pad with zero
+    if embedding.len() < embedding_dim {
+      embedding.push(0.0);
+    }
+
+    // Ensure we have the correct length
+    assert_eq!(embedding.len(), embedding_dim, "Embedding length mismatch");
+
+    // Create tensor from the embedding vector
+    let tensor = Tensor::new(embedding.as_slice(), device)?.to_dtype(dtype)?;
+
+    // Reshape to match expected [1, embedding_dim] shape
+    let result = tensor.reshape((1, embedding_dim))?;
+
+    println!("Final guidance embedding shape: {:?}", result.shape());
+
+    Ok(result)
+  }
+
+  /// Calculate timesteps based on LCM's specific formula
+  /// This is now only used for backward compatibility or custom timesteps
+  pub fn get_timesteps_for_steps(inference_steps: usize) -> Vec<usize> {
+    match inference_steps {
+      // 4 => vec![999, 499, 259, 0],
+      4 => vec![499, 259],
+      8 => vec![999, 879, 759, 639, 499, 379, 259, 139, 0],
+      _ => {
+        let intervals = (999 / (inference_steps - 1)) as usize;
+        let mut timesteps = Vec::with_capacity(inference_steps);
+
+        // First step always starts at 999
+        timesteps.push(999);
+
+        // Generate intermediate steps
+        for i in 1..inference_steps {
+          // The pattern looks approximately linear with a step size determined by the number of steps
+          let ts = 999 - i * intervals;
+          timesteps.push(ts);
+        }
+
+        // Last step is always 0
+        if let Some(last) = timesteps.last_mut() {
+          if *last != 0 {
+            *last = 0;
+          }
+        }
+
+        timesteps
+      },
+    }
+  }
+
+  /// Get the appropriate starting timestep based on strength
+  pub fn get_timestep_for_strength(&self, strength: f64) -> (usize, f64) {
+    // Convert strength to a timestep index
+    let num_steps = self.timesteps.len();
+    let strength_adjusted = (strength * num_steps as f64).floor() as usize;
+    let start_idx = num_steps.saturating_sub(strength_adjusted + 1);
+
+    // Get the timestep at this index
+    let timestep = if start_idx < self.timesteps.len() { self.timesteps[start_idx] } else { *self.timesteps.first().unwrap_or(&999) };
+
+    // Get the corresponding alpha
+    let alpha = if timestep < self.alphas_cumprod.len() { self.alphas_cumprod[timestep] } else { *self.alphas_cumprod.last().unwrap_or(&0.0) };
+
+    println!("For strength {:.2}, using timestep {} (index {}), alpha={:.6}", strength, timestep, start_idx, alpha);
+
+    (timestep, alpha)
+  }
+
+  /// Add noise for img2img based on strength
+  pub fn add_noise_for_img2img(&mut self, original_samples: &Tensor, noise: &Tensor, strength: f64) -> Result<(Tensor, usize)> {
+    // Get starting timestep and alpha
+    let (timestep, alpha) = self.get_timestep_for_strength(strength);
+
+    // Set begin_index based on this timestep
+    let start_idx = self.index_for_timestep(timestep)?;
+    self.set_begin_index(start_idx);
+
+    println!("Adding img2img noise with strength={:.4}, timestep={}, alpha={:.6}", strength, timestep, alpha);
+    println!("  original_samples: shape={:?}, dtype={:?}", original_samples.shape(), original_samples.dtype());
+    println!("  noise: shape={:?}, dtype={:?}", noise.shape(), noise.dtype());
+
+    // Keep operations in original_samples' dtype for now
+    let target_dtype = original_samples.dtype();
+    let noise_converted = noise.to_dtype(target_dtype)?;
+
+    // Calculate coefficients
+    let sqrt_alpha = alpha.sqrt();
+    let sqrt_one_minus_alpha = (1.0 - alpha).sqrt();
+
+    println!("  Using scalar multiplication with target dtype={:?}", target_dtype);
+
+    // Apply noise formula
+    let scaled_samples = original_samples.affine(sqrt_alpha, 0.0)?;
+    let scaled_noise = noise_converted.affine(sqrt_one_minus_alpha, 0.0)?;
+
+    println!("  scaled_samples: shape={:?}, dtype={:?}", scaled_samples.shape(), scaled_samples.dtype());
+    println!("  scaled_noise: shape={:?}, dtype={:?}", scaled_noise.shape(), scaled_noise.dtype());
+
+    // Add scaled noise - should maintain dtype
+    let noised = scaled_samples.add(&scaled_noise)?;
+
+    println!("  noised: shape={:?}, dtype={:?}", noised.shape(), noised.dtype());
+
+    // Return noised image and the timestep to start from
+    Ok((noised, timestep))
+  }
+
+  pub fn index_for_timestep(&self, timestep: usize) -> Result<usize> {
+    // Collect all indices where timestep matches
+    let indices: Vec<usize> = self.timesteps.iter().enumerate().filter_map(|(idx, &ts)| if ts == timestep { Some(idx) } else { None }).collect();
+
+    if indices.is_empty() {
+      return Err(candle::Error::Msg(format!("Timestep {} not found in scheduler timesteps", timestep)));
+    }
+
+    // Take the second index if there are multiple matches, otherwise take the first
+    let pos = if indices.len() > 1 { 1 } else { 0 };
+    Ok(indices[pos])
+  }
+
+  pub fn set_begin_index(&mut self, begin_index: usize) {
+    self.begin_index = Some(begin_index);
+    println!("Set begin_index to {}", begin_index);
+  }
+
+  pub fn init_step_index(&mut self, timestep: usize) -> Result<usize> {
+    self.step_index = if let Some(begin_index) = self.begin_index {
+      // If begin_index is set, use it directly
+      Some(begin_index)
+    } else {
+      // Otherwise find the index based on the timestep
+      // Find the position of timestep in self.timesteps
+      let pos = self.timesteps.iter().position(|&t| t == timestep).ok_or_else(|| candle::Error::Msg(format!("Timestep {} not found in timesteps", timestep)))?;
+      Some(pos)
+    };
+
+    Ok(self.step_index.unwrap())
+  }
+
+  // Updated helper function to match diffusers' get_timesteps logic exactly
+  pub fn retrieve_timesteps(scheduler: &mut LCMScheduler, num_inference_steps: usize, strength: Option<f64>, custom_timesteps: Option<Vec<usize>>) -> Result<Vec<usize>> {
+    // 1. When custom timesteps are provided, use them directly
+    if let Some(custom_ts) = custom_timesteps {
+      println!("Using custom timesteps: {:?}", custom_ts);
+      return Ok(custom_ts);
+    }
+
+    // 2. Get timesteps from scheduler
+    let mut timesteps = scheduler.timesteps().to_vec();
+
+    // 3. For img2img, adjust timesteps based on strength
+    if let Some(strength) = strength {
+      if strength < 1.0 {
+        // Match diffusers logic exactly:
+        // init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+        let init_timestep = ((num_inference_steps as f64 * strength).floor() as usize).min(num_inference_steps);
+
+        // t_start = max(num_inference_steps - init_timestep, 0)
+        let t_start = num_inference_steps.saturating_sub(init_timestep);
+
+        // In diffusers, order is typically 1 for LCM, so we'll simplify here
+        let scheduler_order = 1;
+        let start_idx = t_start * scheduler_order;
+
+        // Slice timesteps from the calculated start position
+        if start_idx < timesteps.len() {
+          println!("Adjusting timesteps for strength {}: starting at index {} (init_timestep={})", strength, start_idx, init_timestep);
+          timesteps = timesteps[start_idx..].to_vec();
+
+          // Set begin_index in the scheduler
+          scheduler.set_begin_index(start_idx);
+        }
+
+        // Return the actual number of inference steps we'll perform
+        let actual_steps = num_inference_steps.saturating_sub(t_start);
+        println!("Will perform {} inference steps due to strength {}", actual_steps, strength);
+      }
+    }
+
+    println!("Final timesteps for inference: {:?}", timesteps);
+    Ok(timesteps)
+  }
+
+  // Add this method to help initialize latents from an image (for img2img)
+  pub fn prepare_img2img_latents(
+    &mut self,
+    image_tensor: &Tensor, // Encoded image from VAE
+    timestep: usize,
+    batch_size: usize,
+    num_images_per_prompt: usize,
+    dtype: DType,
+    device: &Device,
+  ) -> Result<Tensor> {
+    // Store the dtype for consistent operations
+    self.dtype = dtype;
+
+    let latent_dim = image_tensor.dim(2)? / 8; // Assuming standard VAE downsampling
+
+    println!("Preparing img2img latents for timestep={}", timestep);
+    println!("  image_tensor: shape={:?}, dtype={:?}", image_tensor.shape(), image_tensor.dtype());
+
+    // Create noise tensor using the same approach as in main.rs
+    let noise_shape = Shape::from((batch_size * num_images_per_prompt, 4, latent_dim, latent_dim));
+    let noise = Tensor::randn(0f32, 1f32, noise_shape, device)?.to_dtype(dtype)?;
+    println!("  noise: shape={:?}, dtype={:?}", noise.shape(), noise.dtype());
+
+    // Encode image to latents
+    let encoded_latents = image_tensor.clone(); // Assuming already encoded
+
+    // Scale to match VAE expected scale - use affine for scalar multiplication
+    println!("  Using affine for VAE scaling by 0.18215");
+    let scaled_latents = encoded_latents.affine(0.18215, 0.0)?;
+    println!("  scaled_latents: shape={:?}, dtype={:?}", scaled_latents.shape(), scaled_latents.dtype());
+
+    // Add noise based on timestep (this handles the "strength" parameter)
+    let noised_latents = self.add_noise(&scaled_latents, noise, timestep)?;
+    println!("  noised_latents: shape={:?}, dtype={:?}", noised_latents.shape(), noised_latents.dtype());
+
+    Ok(noised_latents)
+  }
+
+  /// Apply classifier-free guidance to model output
+  pub fn apply_guidance(&self, model_output_uncond: &Tensor, model_output_text: &Tensor, guidance_scale: f64) -> Result<Tensor> {
+    // Skip guidance if scale is 1.0 (or very close to it)
+    if (guidance_scale - 1.0).abs() < 1e-5 {
+      return Ok(model_output_text.clone());
+    }
+
+    println!("Applying guidance with scale={:.4}", guidance_scale);
+    println!("  model_output_uncond: shape={:?}, dtype={:?}", model_output_uncond.shape(), model_output_uncond.dtype());
+    println!("  model_output_text: shape={:?}, dtype={:?}", model_output_text.shape(), model_output_text.dtype());
+
+    // Apply the classifier-free guidance formula:
+    // output = uncond + guidance_scale * (text - uncond)
+
+    // First calculate (text - uncond)
+    let guidance_diff = (model_output_text - model_output_uncond)?;
+    println!("  guidance_diff: shape={:?}, dtype={:?}", guidance_diff.shape(), guidance_diff.dtype());
+
+    // Then scale by guidance_scale using affine for scalar multiplication
+    let guidance_factor = guidance_scale - 1.0;
+    println!("  Using affine for scaling by {:.4}", guidance_factor);
+    let scaled_diff = guidance_diff.affine(guidance_factor, 0.0)?;
+    println!("  scaled_diff: shape={:?}, dtype={:?}", scaled_diff.shape(), scaled_diff.dtype());
+
+    // Finally add back to uncond: uncond + guidance_scale * (text - uncond)
+    let guided_output = (model_output_uncond + &scaled_diff)?;
+    println!("  guided_output: shape={:?}, dtype={:?}", guided_output.shape(), guided_output.dtype());
+
+    Ok(guided_output)
+  }
+
+  // Add this helper method to manage dtype consistency across the pipeline
+  pub fn ensure_dtype_consistency(&mut self, input_tensor: &Tensor) -> Result<()> {
+    // Store the current tensor's dtype for reference
+    self.dtype = input_tensor.dtype();
+    println!("Setting LCM scheduler dtype to match input: {:?}", self.dtype);
+    Ok(())
+  }
+
+  // Add a new method to prepare latents with explicit dtype control
+  pub fn prepare_latents(&mut self, batch_size: usize, height: usize, width: usize, num_channels_latents: usize, dtype: DType, device: &Device, generator: Option<u64>) -> Result<Tensor> {
+    // Store the requested dtype
+    self.dtype = dtype;
+
+    // Calculate latent dimensions
+    let latent_height = height / 8;
+    let latent_width = width / 8;
+
+    // Create random latents
+    let shape = Shape::from((batch_size, num_channels_latents, latent_height, latent_width));
+
+    // Note: Candle doesn't have a direct way to set random seeds for tensors
+    // So we'll just use regular random generation for now
+    let latents = if let Some(seed) = generator {
+      println!("Note: Using random latents (seed not supported in this implementation)");
+      // In a real implementation, you would want to set the random seed here
+      Tensor::randn(0f32, 1f32, shape, device)?
+    } else {
+      println!("Generating random latents");
+      Tensor::randn(0f32, 1f32, shape, device)?
+    };
+
+    // Convert to the requested dtype
+    let latents = latents.to_dtype(dtype)?;
+    println!("Generated latents: shape={:?}, dtype={:?}", latents.shape(), latents.dtype());
+
+    // Scale by the init_noise_sigma
+    let scaled_latents = latents.affine(self.init_noise_sigma, 0.0)?;
+
+    Ok(scaled_latents)
+  }
+
+  // Add this method to maintain dtype consistency across operations
+  pub fn match_dtype(&self, tensor: &Tensor, dtype: Option<DType>) -> Result<Tensor> {
+    match dtype {
+      Some(d) => tensor.to_dtype(d),
+      None => Ok(tensor.clone()),
+    }
+  }
+
+  // Add this public getter for the dtype field
+  pub fn dtype(&self) -> DType {
+    self.dtype
+  }
+}
diff --git a/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/mod.rs b/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/mod.rs
index 4c685209cb..d174da34d2 100644
--- a/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/mod.rs
+++ b/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/mod.rs
@@ -41,6 +41,7 @@ pub mod attention;
 pub mod clip;
 pub mod ddim;
 pub mod ddpm;
+pub mod lcm;
 pub mod embeddings;
 pub mod euler_ancestral_discrete;
 pub mod resnet;
@@ -52,8 +53,9 @@ pub mod utils;
 pub mod vae;
 
 use std::sync::Arc;
+use std::collections::HashMap;
 
-use candle::{DType, Device, Result};
+use candle::{DType, Device, Result, Tensor};
 use candle_nn as nn;
 
 use self::schedulers::{Scheduler, SchedulerConfig};
@@ -494,6 +496,28 @@ impl StableDiffusionConfig {
     pub fn build_scheduler(&self, n_steps: usize) -> Result<Box<dyn Scheduler>> {
         self.scheduler.build(n_steps)
     }
+
+    fn unet_config(&self, in_channels: usize) -> Result<unet_2d::UNet2DConditionModelConfig> {
+        Ok(self.unet.clone())
+    }
+
+    pub fn build_unet_with_weights(
+        &self,
+        weights: HashMap<String, Tensor>,
+        device: &Device,
+        in_channels: usize,
+        use_flash_attn: bool,
+        dtype: DType,
+    ) -> Result<unet_2d::UNet2DConditionModel> {
+        let config = self.unet.clone();
+        unet_2d::UNet2DConditionModel::new(
+            nn::VarBuilder::from_tensors(weights, dtype, device),
+            in_channels,
+            4,
+            use_flash_attn,
+            config,
+        )
+    }
 }
 
 pub fn build_clip_transformer<P: AsRef<std::path::Path>>(
diff --git a/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/unet.rs b/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/unet.rs
new file mode 100644
index 0000000000..cb62c4f167
--- /dev/null
+++ b/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/unet.rs
@@ -0,0 +1,27 @@
+impl UNet2DConditionModel {
+    pub fn load_with_weights(
+        config: UNet2DConfig,
+        weights: std::collections::HashMap<String, Tensor>,
+        device: &Device,
+        use_flash_attn: bool,
+        dtype: DType,
+    ) -> Result<Self> {
+        let mut unet = Self::new(&config, use_flash_attn)?;
+        unet.load_tensors_from_hash(weights, device, dtype)?;
+        Ok(unet)
+    }
+
+    fn load_tensors_from_hash(
+        &mut self,
+        mut weights: std::collections::HashMap<String, Tensor>,
+        device: &Device,
+        dtype: DType,
+    ) -> Result<()> {
+        // Load the weights from the HashMap into the model
+        for (name, tensor) in weights.drain() {
+            let tensor = tensor.to_device(device)?.to_dtype(dtype)?;
+            self.load_tensor(&name, tensor)?;
+        }
+        Ok(())
+    }
+} 
\ No newline at end of file
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 cbef3316dd..6215b5efda 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,397 +5,301 @@
 use super::embeddings::{TimestepEmbedding, Timesteps};
 use super::unet_2d_blocks::*;
 use crate::models::with_tracing::{conv2d, Conv2d};
-use candle::{Result, Tensor};
+use candle::{Result, Tensor, Device, DType};
 use candle_nn as nn;
 use candle_nn::Module;
+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 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,
 }
 
 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,
-        }
-    }
+  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 }
+  }
 }
 
 #[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 = b_channels * 4;
-        let conv_cfg = nn::Conv2dConfig {
-            padding: 1,
-            ..Default::default()
-        };
-        let conv_in = conv2d(in_channels, b_channels, 3, conv_cfg, vs.pp("conv_in"))?;
+  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 = b_channels * 4;
+    let conv_cfg = nn::Conv2dConfig { padding: 1, ..Default::default() };
+    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 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];
+    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,
-                };
-
-                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()
+        // 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 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)?
+        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 {
-            xs.clone()
-        };
-        // 1. time
-        let emb = (Tensor::ones(bsize, xs.dtype(), device)? * timestep)?;
-        let emb = self.time_proj.forward(&emb)?;
-        let 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(&emb))?,
-                UNetDownBlock::CrossAttn(b) => {
-                    b.forward(&xs, Some(&emb), Some(encoder_hidden_states))?
-                }
-            };
-            down_block_res_xs.extend(res_xs);
-            xs = _xs;
+          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<_>>>()?;
 
-        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;
+    // 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() };
 
-        // 4. mid
-        let xs = self
-            .mid_block
-            .forward(&xs, Some(&emb), Some(encoder_hidden_states))?;
-        let xs = match mid_block_additional_residual {
-            None => xs,
-            Some(m) => (m + xs)?,
+    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,
         };
-        // 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(&emb), upsample_size)?,
-                UNetUpBlock::CrossAttn(b) => b.forward(
-                    &xs,
-                    &res_xs,
-                    Some(&emb),
-                    upsample_size,
-                    Some(encoder_hidden_states),
-                )?,
-            };
+
+        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))
         }
-        // 6. post-process
-        let xs = self.conv_norm_out.forward(&xs)?;
-        let xs = nn::ops::silu(&xs)?;
-        self.conv_out.forward(&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 {
+      // Print shapes for debugging
+      println!("Guidance embedding shape: {:?}", guidance_emb.shape());
+      println!("Time embedding shape: {:?}", time_emb.shape());
+
+      // Get the batch size and embedding dimension
+      let (batch_size, emb_dim) = time_emb.dims2()?;
+
+      // Reshape guidance embedding to match time embedding dimensions
+      let guidance_emb = guidance_emb.reshape(&[1, emb_dim])?.expand(&[batch_size, emb_dim])?;
+
+      // Add to time embeddings
+      time_emb = (time_emb + guidance_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;
+    }
+
+    // 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)
+  }
 }