loading indicator spinner update

added signals for the loading indicator

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"

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)
+}

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
+  }
+}

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>>(

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(())
+    }
+} 

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)
+  }
 }

frontend/realtime/src/KonvaRootComponent/SignaledLoadingIndicator.tsx

@@ -0,0 +1,54 @@
+import { useSignalEffect } from "@preact/signals-react";
+import { LoadingBar, LoadingBarStatus } from "~/components/ui";
+import {
+  isLoadingVisible,
+  loadingProgress,
+} from "~/signals/uiEvents/loadingIndicator";
+
+export const SignaledLoadingIndicator = () => {
+  // This makes component re-render whenever the signals change
+  useSignalEffect(() => {
+    isLoadingVisible.value;
+    loadingProgress.value;
+  });
+
+  if (!isLoadingVisible.value) {
+    return null;
+  }
+
+  return (
+    <div className="flex min-w-[328px] flex-col items-center">
+      <LoadingBar
+        progress={loadingProgress.value}
+        status={LoadingBarStatus.LOADING}
+      />
+      <div className="glass flex w-fit items-center gap-1.5 rounded-md px-2 py-1 text-xs font-medium opacity-60">
+        <svg
+          className="h-3 w-3 animate-spin text-white/40"
+          viewBox="0 0 64 64"
+          fill="none"
+          xmlns="http://www.w3.org/2000/svg"
+          width="24"
+          height="24"
+        >
+          <path
+            d="M32 3C35.8083 3 39.5794 3.75011 43.0978 5.20749C46.6163 6.66488 49.8132 8.80101 52.5061 11.4939C55.199 14.1868 57.3351 17.3837 58.7925 20.9022C60.2499 24.4206 61 28.1917 61 32C61 35.8083 60.2499 39.5794 58.7925 43.0978C57.3351 46.6163 55.199 49.8132 52.5061 52.5061C49.8132 55.199 46.6163 57.3351 43.0978 58.7925C39.5794 60.2499 35.8083 61 32 61C28.1917 61 24.4206 60.2499 20.9022 58.7925C17.3837 57.3351 14.1868 55.199 11.4939 52.5061C8.801 49.8132 6.66487 46.6163 5.20749 43.0978C3.7501 39.5794 3 35.8083 3 32C3 28.1917 3.75011 24.4206 5.2075 20.9022C6.66489 17.3837 8.80101 14.1868 11.4939 11.4939C14.1868 8.80099 17.3838 6.66487 20.9022 5.20749C24.4206 3.7501 28.1917 3 32 3L32 3Z"
+            stroke="currentColor"
+            strokeWidth="5"
+            strokeLinecap="round"
+            strokeLinejoin="round"
+          ></path>
+          <path
+            d="M32 3C36.5778 3 41.0906 4.08374 45.1692 6.16256C49.2477 8.24138 52.7762 11.2562 55.466 14.9605C58.1558 18.6647 59.9304 22.9531 60.6448 27.4748C61.3591 31.9965 60.9928 36.6232 59.5759 40.9762"
+            stroke="currentColor"
+            strokeWidth="5"
+            strokeLinecap="round"
+            strokeLinejoin="round"
+            className="text-white/90"
+          ></path>
+        </svg>
+        Downloading models - {loadingProgress.value}%
+      </div>
+    </div>
+  );
+};

frontend/realtime/src/KonvaRootComponent/SignaledModelButtons.tsx

@@ -0,0 +1,45 @@
+import { twMerge } from "tailwind-merge";
+import { ModelButton } from "~/components/features/ModelSelection/ModelButton";
+import { useSignalEffect } from "@preact/signals-react";
+import { selectedModel, selectedLora } from "~/signals/uiEvents/modelSelection";
+import { getModelById, getLoraById } from "~/data/models";
+
+export const SignaledModelButtons = () => {
+  // This makes the component re-render when the selected model or lora changes
+  useSignalEffect(() => {
+    selectedModel.value;
+    selectedLora.value;
+  });
+
+  // Get the selected model's image or use default
+  const selectedModelData = selectedModel.value
+    ? getModelById(selectedModel.value)
+    : null;
+  const modelIconSrc = selectedModelData?.imageUrl || "";
+
+  // Get the selected LoRA's image or use default
+  const selectedLoraData = selectedLora.value
+    ? getLoraById(selectedLora.value)
+    : null;
+  const loraIconSrc = selectedLoraData?.imageUrl || "";
+
+  return (
+    <div className={twMerge("z-20 flex h-fit w-fit items-center gap-2")}>
+      <ModelButton
+        iconSrc={modelIconSrc}
+        iconAlt={selectedModelData?.name || "Model"}
+        label={selectedModelData?.name || "Realistic"}
+        subtitle="Model"
+        type="model"
+      />
+
+      <ModelButton
+        iconSrc={loraIconSrc}
+        iconAlt={selectedLoraData?.name || "LoRA"}
+        label={selectedLoraData?.name || "Detailed Tweaker"}
+        subtitle="LoRA"
+        type="lora"
+      />
+    </div>
+  );
+};

frontend/realtime/src/components/features/ModelSelection/LoraSelectionDialog.tsx

@@ -0,0 +1,64 @@
+import { DialogTitle } from "@headlessui/react";
+import { useSignalEffect } from "@preact/signals-react";
+import { ModelCard } from "./ModelCard";
+import { loras, isLoraCompatibleWithModel } from "~/data/models";
+import {
+  selectedModel,
+  selectedLora,
+  dispatchers,
+} from "~/signals/uiEvents/modelSelection";
+
+interface LoraSelectionDialogProps {
+  onClose: () => void;
+}
+
+export const LoraSelectionDialog = ({ onClose }: LoraSelectionDialogProps) => {
+  // This makes component re-render whenever signals change
+  useSignalEffect(() => {
+    selectedModel.value;
+    selectedLora.value;
+  });
+
+  const handleSelectLora = (loraId: string) => {
+    dispatchers.setSelectedLora(loraId);
+    onClose(); // Close the dialog after selection
+  };
+
+  const currentModelId = selectedModel.value;
+
+  return (
+    <div className="flex flex-col gap-4">
+      <div>
+        <DialogTitle className="text-2xl font-bold">Select LoRA</DialogTitle>
+        <p className="text-sm text-gray-400">
+          Pick a LoRA to enhance your selected model
+          {!currentModelId && (
+            <span className="ml-1 text-yellow-500">
+              (Select a base model first)
+            </span>
+          )}
+        </p>
+      </div>
+
+      <div className="grid grid-cols-4 gap-4">
+        {loras.map((lora) => (
+          <ModelCard
+            key={lora.id}
+            id={lora.id}
+            name={lora.name}
+            imageUrl={lora.imageUrl}
+            isDownloaded={lora.isDownloaded}
+            isSelected={selectedLora.value === lora.id}
+            isCompatible={
+              currentModelId
+                ? isLoraCompatibleWithModel(lora.id, currentModelId)
+                : false
+            }
+            type="lora"
+            onSelect={handleSelectLora}
+          />
+        ))}
+      </div>
+    </div>
+  );
+};

frontend/realtime/src/components/features/ModelSelection/ModelButton.tsx

@@ -0,0 +1,93 @@
+import { useState } from "react";
+import { Button } from "~/components/ui/Button";
+import { BaseDialog } from "~/components/ui/BaseDialog";
+import { useSignalEffect } from "@preact/signals-react";
+import { selectedModel, selectedLora } from "~/signals/uiEvents/modelSelection";
+import { getModelById, getLoraById } from "~/data/models";
+import { ModelSelectionDialog } from "./ModelSelectionDialog";
+import { LoraSelectionDialog } from "./LoraSelectionDialog";
+
+interface ModelButtonProps {
+  iconSrc: string;
+  iconAlt: string;
+  label?: string;
+  subtitle?: string;
+  type: "model" | "lora";
+}
+
+export const ModelButton = ({
+  iconSrc,
+  iconAlt,
+  label,
+  subtitle,
+  type,
+}: ModelButtonProps) => {
+  const [isOpen, setIsOpen] = useState(false);
+
+  // Update button label based on selection
+  useSignalEffect(() => {
+    // This will re-render component when these signals change
+    selectedModel.value;
+    selectedLora.value;
+  });
+
+  // Set button text based on current selection
+  let buttonLabel = label;
+  let buttonSubtitle = subtitle;
+
+  if (type === "model" && selectedModel.value) {
+    const model = getModelById(selectedModel.value);
+    if (model) {
+      buttonLabel = model.name;
+      buttonSubtitle = "Model";
+    }
+  }
+
+  if (type === "lora" && selectedLora.value) {
+    const lora = getLoraById(selectedLora.value);
+    if (lora) {
+      buttonLabel = lora.name;
+      buttonSubtitle = "LoRA";
+    }
+  }
+
+  return (
+    <>
+      <Button
+        variant="secondary"
+        className="glass flex w-40 items-center justify-start gap-2 bg-ui-controls/70 p-1"
+        onClick={() => setIsOpen(true)}
+      >
+        <img
+          src={iconSrc}
+          alt={iconAlt}
+          className="h-8 w-8 rounded-md bg-white/10"
+        />
+        {buttonSubtitle ? (
+          <div className="flex flex-col gap-0 truncate text-start">
+            <span className="truncate text-[11px] opacity-60">
+              {buttonSubtitle}
+            </span>
+            <span className="truncate text-[13px] opacity-80">
+              {buttonLabel}
+            </span>
+          </div>
+        ) : (
+          <span className="truncate">{buttonLabel}</span>
+        )}
+      </Button>
+
+      <BaseDialog
+        isOpen={isOpen}
+        onClose={() => setIsOpen(false)}
+        className="max-w-3xl"
+      >
+        {type === "model" ? (
+          <ModelSelectionDialog onClose={() => setIsOpen(false)} />
+        ) : (
+          <LoraSelectionDialog onClose={() => setIsOpen(false)} />
+        )}
+      </BaseDialog>
+    </>
+  );
+};

frontend/realtime/src/components/features/ModelSelection/ModelCard.tsx

@@ -0,0 +1,156 @@
+import { faDownToLine } from "@fortawesome/pro-solid-svg-icons";
+import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
+import { twMerge } from "tailwind-merge";
+import { useSignalEffect } from "@preact/signals-react";
+import { useState } from "react";
+
+import {
+  downloadingItems,
+  dispatchers,
+} from "~/signals/uiEvents/modelSelection";
+import { downloadModel } from "~/services/modelDownloadService";
+
+interface ModelCardProps {
+  id: string;
+  name: string;
+  imageUrl: string;
+  isDownloaded: boolean;
+  isSelected: boolean;
+  isCompatible: boolean;
+  type: "model" | "lora";
+  onSelect: (id: string) => void;
+}
+
+export const ModelCard = ({
+  id,
+  name,
+  imageUrl,
+  isDownloaded,
+  isSelected,
+  isCompatible,
+  type,
+  onSelect,
+}: ModelCardProps) => {
+  // Local state to track download status and progress
+  const [progress, setProgress] = useState(0);
+  const [isDownloading, setIsDownloading] = useState(false);
+  const [localIsDownloaded, setLocalIsDownloaded] = useState(isDownloaded);
+
+  // Update local state when downloadingItems changes
+  useSignalEffect(() => {
+    const downloads = downloadingItems.value;
+    const isCurrentlyDownloading = !!downloads[id];
+    setIsDownloading(isCurrentlyDownloading);
+
+    if (isCurrentlyDownloading) {
+      setProgress(downloads[id].progress);
+    }
+  });
+
+  // Function to handle the download
+  const handleDownload = async () => {
+    try {
+      // Use the download service to start the download
+      await downloadModel(id, type);
+
+      // For the dummy implementation, we'll update the local state when download completes
+      const checkDownloadStatus = setInterval(() => {
+        if (!dispatchers.isDownloading(id)) {
+          clearInterval(checkDownloadStatus);
+          setLocalIsDownloaded(true);
+        }
+      }, 500);
+    } catch (error) {
+      console.error(`Error downloading ${type} ${id}:`, error);
+    }
+  };
+
+  const handleClick = () => {
+    if (!isCompatible) return; // Do nothing if not compatible
+
+    if (isDownloading) return; // Do nothing if already downloading
+
+    if (localIsDownloaded) {
+      onSelect(id);
+    } else {
+      handleDownload();
+    }
+  };
+
+  return (
+    <div
+      className={twMerge(
+        "relative flex aspect-square flex-col items-center overflow-hidden rounded-xl border-2 bg-white/5 transition-all",
+        isSelected ? "border-primary" : "border-transparent",
+        !isCompatible ? "opacity-50 grayscale" : "hover:border-primary",
+        "cursor-pointer",
+      )}
+      onClick={handleClick}
+    >
+      {/* Selected indicator */}
+      {isSelected && (
+        <div className="absolute right-2 top-2 z-10 flex h-6 w-6 items-center justify-center rounded-full bg-primary text-white">
+          <svg
+            xmlns="http://www.w3.org/2000/svg"
+            width="16"
+            height="16"
+            viewBox="0 0 24 24"
+            fill="none"
+            stroke="currentColor"
+            strokeWidth="3"
+            strokeLinecap="round"
+            strokeLinejoin="round"
+          >
+            <polyline points="20 6 9 17 4 12"></polyline>
+          </svg>
+        </div>
+      )}
+
+      {/* Image container */}
+      <div className="relative h-full w-full overflow-hidden">
+        {/* Downloading indicator */}
+        {isDownloading && (
+          <div className="absolute inset-0 flex h-full w-full flex-col items-center justify-center bg-gray-900/75">
+            <svg
+              className="h-8 w-8 animate-spin text-white/40"
+              viewBox="0 0 64 64"
+              fill="none"
+              xmlns="http://www.w3.org/2000/svg"
+              width="24"
+              height="24"
+            >
+              <path
+                d="M32 3C35.8083 3 39.5794 3.75011 43.0978 5.20749C46.6163 6.66488 49.8132 8.80101 52.5061 11.4939C55.199 14.1868 57.3351 17.3837 58.7925 20.9022C60.2499 24.4206 61 28.1917 61 32C61 35.8083 60.2499 39.5794 58.7925 43.0978C57.3351 46.6163 55.199 49.8132 52.5061 52.5061C49.8132 55.199 46.6163 57.3351 43.0978 58.7925C39.5794 60.2499 35.8083 61 32 61C28.1917 61 24.4206 60.2499 20.9022 58.7925C17.3837 57.3351 14.1868 55.199 11.4939 52.5061C8.801 49.8132 6.66487 46.6163 5.20749 43.0978C3.7501 39.5794 3 35.8083 3 32C3 28.1917 3.75011 24.4206 5.2075 20.9022C6.66489 17.3837 8.80101 14.1868 11.4939 11.4939C14.1868 8.80099 17.3838 6.66487 20.9022 5.20749C24.4206 3.7501 28.1917 3 32 3L32 3Z"
+                stroke="currentColor"
+                strokeWidth="5"
+                strokeLinecap="round"
+                strokeLinejoin="round"
+              ></path>
+              <path
+                d="M32 3C36.5778 3 41.0906 4.08374 45.1692 6.16256C49.2477 8.24138 52.7762 11.2562 55.466 14.9605C58.1558 18.6647 59.9304 22.9531 60.6448 27.4748C61.3591 31.9965 60.9928 36.6232 59.5759 40.9762"
+                stroke="currentColor"
+                strokeWidth="5"
+                strokeLinecap="round"
+                strokeLinejoin="round"
+                className="text-white/90"
+              ></path>
+            </svg>
+            <div className="mt-2 text-center font-medium text-white">
+              {progress}%
+            </div>
+          </div>
+        )}
+
+        <img src={imageUrl} alt={name} className="h-full w-full object-cover" />
+
+        {/* Title */}
+        <div className="absolute bottom-0 mt-2 flex w-full items-center justify-between bg-[#3E3E41]/70 px-2 py-1 text-sm font-medium backdrop-blur-[2px]">
+          {name}
+          {!localIsDownloaded && !isDownloading && (
+            <FontAwesomeIcon icon={faDownToLine} className="text-white" />
+          )}
+        </div>
+      </div>
+    </div>
+  );
+};

frontend/realtime/src/components/features/ModelSelection/ModelSelectionDialog.tsx

@@ -0,0 +1,56 @@
+import { DialogTitle } from "@headlessui/react";
+import { useSignalEffect } from "@preact/signals-react";
+import { ModelCard } from "./ModelCard";
+import { models } from "~/data/models";
+import { selectedModel, dispatchers } from "~/signals/uiEvents/modelSelection";
+import { useEffect } from "react";
+
+interface ModelSelectionDialogProps {
+  onClose: () => void;
+}
+
+export const ModelSelectionDialog = ({
+  onClose,
+}: ModelSelectionDialogProps) => {
+  // This makes component re-render whenever signals change
+  useSignalEffect(() => {
+    selectedModel.value;
+  });
+
+  // Ensure a model is selected when the dialog opens
+  useEffect(() => {
+    if (selectedModel.value === null && models.length > 0) {
+      dispatchers.setSelectedModel(models[0].id);
+    }
+  }, []);
+
+  const handleSelectModel = (modelId: string) => {
+    dispatchers.setSelectedModel(modelId);
+    onClose(); // Close the dialog after selection
+  };
+
+  return (
+    <div className="flex flex-col gap-4">
+      <div>
+        <DialogTitle className="text-2xl font-bold">Select Model</DialogTitle>
+        <p className="text-sm text-gray-400">Pick your base generation model</p>
+      </div>
+
+      <div className="grid grid-cols-4 gap-4">
+        {models.map((model) => (
+          <ModelCard
+            key={model.id}
+            id={model.id}
+            name={model.name}
+            imageUrl={model.imageUrl}
+            isDownloaded={model.isDownloaded}
+            isSelected={selectedModel.value === model.id}
+            isCompatible={true} // Base models are always compatible
+            type="model"
+            onSelect={handleSelectModel}
+          />
+        ))}
+      </div>
+    </div>
+  );
+};

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

@@ -1,89 +1,14 @@
-import { twMerge } from "tailwind-merge";
-import { Button } from "~/components/ui/Button";
-import { useState } from "react";
-import { BaseDialog } from "~/components/ui/BaseDialog";
-import { DialogTitle } from "@headlessui/react";
-
-interface ModelButtonProps {
-  iconSrc: string;
-  iconAlt: string;
-  label?: string;
-  subtitle?: string;
-  children?: React.ReactNode;
-}
-
-const ModelButton = ({
-  iconSrc,
-  iconAlt,
-  label,
-  subtitle,
-  children,
-}: ModelButtonProps) => {
-  const [isOpen, setIsOpen] = useState(false);
-
-  return (
-    <>
-      <Button
-        variant="secondary"
-        className="glass flex w-40 items-center justify-start gap-2 bg-ui-controls/70 p-1"
-        onClick={() => setIsOpen(true)}
-      >
-        <img
-          src={iconSrc}
-          alt={iconAlt}
-          className="h-8 w-8 rounded-md bg-white/10"
-        />
-        {subtitle ? (
-          <div className="flex flex-col gap-0 truncate text-start">
-            <span className="truncate text-[11px] opacity-60">{subtitle}</span>
-            <span className="truncate text-[13px] opacity-80">{label}</span>
-          </div>
-        ) : (
-          <span className="truncate">{label}</span>
-        )}
-      </Button>
-
-      <BaseDialog isOpen={isOpen} onClose={() => setIsOpen(false)}>
-        {children}
-      </BaseDialog>
-    </>
-  );
-};
+import { SignaledLoadingIndicator } from "~/KonvaRootComponent/SignaledLoadingIndicator";
+import { SignaledModelButtons } from "~/KonvaRootComponent/SignaledModelButtons";
 
 export const ToolbarTopCenter = () => {
   return (
-    <>
-      {import.meta.env.DEV && ( // only show in dev mode for now
-        <div className={twMerge("z-20 flex h-fit w-fit items-center gap-2")}>
-          <ModelButton
-            iconSrc="/resources/icons/model.png"
-            iconAlt="Model"
-            label="Realistic"
-            subtitle="Model"
-          >
-            <div className="flex flex-col gap-2 text-sm">
-              <DialogTitle className="text-2xl font-bold">
-                Select Model
-              </DialogTitle>
-              <p>Model configuration options will go here</p>
-            </div>
-          </ModelButton>
+    <div className="flex flex-col gap-3">
+      {/* Model selection buttons (only shown in dev mode for now) */}
+      {import.meta.env.DEV && <SignaledModelButtons />}
 
-          <ModelButton
-            iconSrc="/resources/icons/lora.png"
-            iconAlt="LoRA"
-            label="Detailed Tweaker"
-            subtitle="LoRA"
-          >
-            <div className="flex flex-col gap-2 text-sm">
-              <DialogTitle className="text-2xl font-bold">
-                Select LoRA
-              </DialogTitle>
-              <p>LoRA configuration options will go here</p>
-            </div>
-          </ModelButton>
-        </div>
-      )}
-    </>
+      {/* Loading models indicator (only shown in dev mode for now) */}
+      {import.meta.env.DEV && <SignaledLoadingIndicator />}
+    </div>
   );
 };

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

@@ -35,17 +35,14 @@ export const LoadingBar = ({
     >
       <div
         className={twMerge(
-          "h-2.5 w-full rounded-full",
+          "h-1.5 w-full rounded-full",
           status && status === LoadingBarStatus.IDLE
-            ? "border-1 bg-gray-50 ring-1 ring-inset ring-gray-200"
-            : "bg-gray-200",
+            ? "bg-gray-50"
+            : "pointer-events-none border-0 bg-gray-600",
         )}
       >
         <div
-          className={twMerge(
-            "h-2.5 rounded-full bg-primary-500",
-            status === LoadingBarStatus.LOADING && "animate-pulse",
-          )}
+          className={twMerge("h-full rounded-full bg-primary/80")}
           style={{ width: `${progress}%` }}
         />
       </div>

frontend/realtime/src/data/models.ts

@@ -0,0 +1,143 @@
+// Model and LoRA type definitions
+export interface ModelData {
+  id: string;
+  name: string;
+  type: string;
+  imageUrl: string;
+  description?: string;
+  isDownloaded: boolean;
+  compatibleLoraIds?: string[]; // LoRAs that work with this model
+}
+
+export interface LoraData {
+  id: string;
+  name: string;
+  imageUrl: string;
+  description?: string;
+  isDownloaded: boolean;
+  compatibleModelIds: string[]; // Models this LoRA works with
+}
+
+// Available models dummy data
+export const models: ModelData[] = [
+  {
+    id: "realistic",
+    name: "Realistic",
+    type: "Realistic",
+    imageUrl: "/images/models/realistic.png",
+    description: "Photorealistic image generation",
+    isDownloaded: true,
+    compatibleLoraIds: [
+      "detailed_tweaker",
+      "scenery_enhancer",
+      "portrait_plus",
+    ],
+  },
+  {
+    id: "flat_2d",
+    name: "Flat 2D",
+    type: "Flat 2D",
+    imageUrl: "/images/models/flat_2d.png",
+    description: "Flat 2D illustration style",
+    isDownloaded: true,
+    compatibleLoraIds: ["toon_enhancer", "detailed_tweaker"],
+  },
+  {
+    id: "ghibli",
+    name: "Ghibli",
+    type: "Ghibli",
+    imageUrl: "/images/models/ghibli.png",
+    description: "Studio Ghibli inspired style",
+    isDownloaded: false,
+    compatibleLoraIds: ["anime_style"],
+  },
+  {
+    id: "scenery",
+    name: "Scenery",
+    type: "Scenery",
+    imageUrl: "/images/models/scenery.png",
+    description: "Beautiful landscape generation",
+    isDownloaded: true,
+    compatibleLoraIds: ["scenery_enhancer"],
+  },
+  {
+    id: "concept",
+    name: "Concept",
+    type: "Concept",
+    imageUrl: "/images/models/concept.png",
+    description: "Abstract concept art generation",
+    isDownloaded: false,
+    compatibleLoraIds: ["detailed_tweaker", "scenery_enhancer"],
+  },
+];
+
+// Available LoRAs
+export const loras: LoraData[] = [
+  {
+    id: "detailed_tweaker",
+    name: "Detailed Tweaker",
+    imageUrl: "/resources/loras/detailedtweaker.jpg",
+    description: "Enhances fine details in generated images",
+    isDownloaded: true,
+    compatibleModelIds: [
+      "realistic",
+      "concept",
+      "anime",
+      "flat_2d",
+      "flux_realtime",
+    ],
+  },
+  {
+    id: "portrait_plus",
+    name: "Portrait+",
+    imageUrl: "/resources/loras/portraitplus.jpg",
+    description: "Improves portrait generation",
+    isDownloaded: false,
+    compatibleModelIds: ["realistic", "flux_realtime"],
+  },
+  {
+    id: "scenery_enhancer",
+    name: "Scenery Enhancer",
+    imageUrl: "/resources/loras/sceneryenhancer.jpg",
+    description: "Enhances landscape generation",
+    isDownloaded: true,
+    compatibleModelIds: ["realistic", "concept", "scenery"],
+  },
+  {
+    id: "anime_style",
+    name: "Anime Style+",
+    imageUrl: "/resources/loras/animestyle.jpg",
+    description: "Enhances anime art style",
+    isDownloaded: false,
+    compatibleModelIds: ["anime", "ghibli"],
+  },
+  {
+    id: "toon_enhancer",
+    name: "Toon Enhancer",
+    imageUrl: "/resources/loras/toonenhancer.jpg",
+    description: "Enhances cartoon styles",
+    isDownloaded: true,
+    compatibleModelIds: ["cartoon", "flat_2d"],
+  },
+];
+
+// Helper functions
+export const getModelById = (id: string): ModelData | undefined => {
+  return models.find((model) => model.id === id);
+};
+
+export const getLoraById = (id: string): LoraData | undefined => {
+  return loras.find((lora) => lora.id === id);
+};
+
+export const isLoraCompatibleWithModel = (
+  loraId: string,
+  modelId: string,
+): boolean => {
+  const lora = getLoraById(loraId);
+  return lora ? lora.compatibleModelIds.includes(modelId) : false;
+};
+
+export const getCompatibleLorasForModel = (modelId: string): LoraData[] => {
+  return loras.filter((lora) => lora.compatibleModelIds.includes(modelId));
+};

frontend/realtime/src/services/modelDownloadService.ts

@@ -0,0 +1,116 @@
+import { dispatchers } from "~/signals/uiEvents/modelSelection";
+import { models, loras, ModelData, LoraData } from "~/data/models";
+
+//THIS IS A VERY DUMMY DOWNLOAD IMPLEMENTATION AND SHOULD BE REPLACED WITH A REAL ONE FROM THE BACKEND OR SOMETHING BECAUSE IM NOT SURE - BFlat
+
+// // Interface for download progress events from Tauri backend
+// interface DownloadProgressEvent {
+//   id: string;
+//   progress: number;
+//   type: "model" | "lora";
+// }
+
+// // Interface for download completion events from Tauri backend
+// interface DownloadCompleteEvent {
+//   id: string;
+//   type: "model" | "lora";
+//   success: boolean;
+//   error?: string;
+// }
+
+// Initialize listeners for download events (dummy implementation)
+export const initializeDownloadListeners = async () => {
+  console.log("Download listeners initialized (dummy implementation)");
+  // No actual listeners in the dummy implementation
+  return Promise.resolve();
+};
+
+// Start downloading a model or lora (dummy implementation)
+export const downloadModel = async (id: string, type: "model" | "lora") => {
+  try {
+    // Start tracking the download in our state
+    dispatchers.startDownload(id, type);
+
+    // Simulate download progress
+    simulateDownloadProgress(id, type);
+
+    return true;
+  } catch (error) {
+    console.error(`Failed to start download for ${type} ${id}:`, error);
+    dispatchers.finishDownload(id);
+    return false;
+  }
+};
+
+// Simulate download progress
+const simulateDownloadProgress = (id: string, type: "model" | "lora") => {
+  let progress = 0;
+
+  // Create an interval to update progress
+  const interval = setInterval(() => {
+    progress += 5;
+
+    // Update progress in state
+    dispatchers.updateDownloadProgress(id, progress);
+
+    // When download completes
+    if (progress >= 100) {
+      clearInterval(interval);
+
+      // Finish the download
+      setTimeout(() => {
+        dispatchers.finishDownload(id);
+
+        // Update the local data to mark the model/lora as downloaded
+        updateDownloadedState(id, type, true);
+
+        console.log(`Download completed for ${type} ${id}`);
+      }, 500);
+    }
+  }, 200); // Update every 200ms
+};
+
+// Update the downloaded state in our local data
+const updateDownloadedState = (
+  id: string,
+  type: "model" | "lora",
+  isDownloaded: boolean,
+) => {
+  if (type === "model") {
+    // Find and update the model
+    const modelIndex = models.findIndex((model) => model.id === id);
+    if (modelIndex >= 0) {
+      const updatedModel: ModelData = {
+        ...models[modelIndex],
+        isDownloaded,
+      };
+
+      // Update the models array
+      models[modelIndex] = updatedModel;
+    }
+  } else {
+    // Find and update the lora
+    const loraIndex = loras.findIndex((lora) => lora.id === id);
+    if (loraIndex >= 0) {
+      const updatedLora: LoraData = {
+        ...loras[loraIndex],
+        isDownloaded,
+      };
+
+      // Update the loras array
+      loras[loraIndex] = updatedLora;
+    }
+  }
+};
+
+// Get all available models (dummy implementation)
+export const fetchAvailableModels = async (): Promise<ModelData[]> => {
+  // Just return the models from the data file
+  return Promise.resolve([...models]);
+};
+
+// Get all available loras (dummy implementation)
+export const fetchAvailableLoras = async (): Promise<LoraData[]> => {
+  // Just return the loras from the data file
+  return Promise.resolve([...loras]);
+};

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

@@ -28,6 +28,14 @@ import {
   dispatchers as promptDispatchers,
   events as promptEvents,
 } from "./promptSettings";
+import {
+  dispatchers as loadingIndicatorDispatchers,
+  events as loadingIndicatorEvents,
+} from "./loadingIndicator";
+import {
+  dispatchers as modelSelectionDispatchers,
+  events as modelSelectionEvents,
+} from "./modelSelection";
 
 export const uiEvents = {
   ...addMediaEvents,
@@ -39,6 +47,8 @@ export const uiEvents = {
   toolbarNode: toolbarNodeEvents,
   toolbarVideoExtraction: toolbarVideoExtractionEvents,
   promptEvents: promptEvents,
+  loadingIndicator: loadingIndicatorEvents,
+  modelSelection: modelSelectionEvents,
 };
 export const dispatchUiEvents = {
   ...addMediaDispatchers,
@@ -50,4 +60,6 @@ export const dispatchUiEvents = {
   toolbarNode: toolbarNodeDispatchers,
   toolbarVideoExtraction: toolbarVideoExtractionDispatchers,
   promptSettings: promptDispatchers,
+  loadingIndicator: loadingIndicatorDispatchers,
+  modelSelection: modelSelectionDispatchers,
 };

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

@@ -0,0 +1,40 @@
+import { effect, signal } from "@preact/signals-react";
+
+// Model loading settings
+// - Is loading visible
+// - Loading progress percentage
+export const DEFAULT_IS_LOADING_VISIBLE = false;
+export const isLoadingVisible = signal<boolean>(DEFAULT_IS_LOADING_VISIBLE);
+const setIsLoadingVisible = (visible: boolean) => {
+  isLoadingVisible.value = visible;
+};
+
+const onIsLoadingVisibleChanged = (callback: (visible: boolean) => void) => {
+  effect(() => {
+    callback(isLoadingVisible.value);
+  });
+};
+
+export const DEFAULT_LOADING_PROGRESS = 0;
+export const loadingProgress = signal<number>(DEFAULT_LOADING_PROGRESS);
+const setLoadingProgress = (progress: number) => {
+  // Round to integer, then clamp between 0 and 100
+  loadingProgress.value = Math.max(0, Math.min(100, Math.round(progress)));
+};
+
+const onLoadingProgressChanged = (callback: (progress: number) => void) => {
+  effect(() => {
+    callback(loadingProgress.value);
+  });
+};
+
+// EXPORTS
+export const dispatchers = {
+  setIsLoadingVisible,
+  setLoadingProgress,
+};
+
+export const events = {
+  onIsLoadingVisibleChanged,
+  onLoadingProgressChanged,
+};

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

@@ -0,0 +1,100 @@
+import { effect, signal } from "@preact/signals-react";
+import { models } from "~/data/models";
+
+// Define the signals for selected models and download states
+export const selectedModel = signal<string | null>(
+  models.length > 0 ? models[0].id : null,
+);
+export const selectedLora = signal<string | null>(null);
+
+// Replace single download tracking with a map to track multiple downloads
+export interface DownloadInfo {
+  progress: number;
+  type: "model" | "lora";
+}
+
+// Map of modelId/loraId to download progress
+export const downloadingItems = signal<Record<string, DownloadInfo>>({});
+
+// Dispatchers
+const setSelectedModel = (modelId: string | null) => {
+  selectedModel.value = modelId;
+};
+
+const setSelectedLora = (loraId: string | null) => {
+  selectedLora.value = loraId;
+};
+
+// Start downloading a model or lora
+const startDownload = (id: string, type: "model" | "lora") => {
+  const newDownloads = { ...downloadingItems.value };
+  newDownloads[id] = { progress: 0, type };
+  downloadingItems.value = newDownloads;
+};
+
+// Update the progress of a downloading item
+const updateDownloadProgress = (id: string, progress: number) => {
+  if (!downloadingItems.value[id]) return;
+
+  const newDownloads = { ...downloadingItems.value };
+  newDownloads[id] = {
+    ...newDownloads[id],
+    progress: Math.max(0, Math.min(100, Math.round(progress))),
+  };
+  downloadingItems.value = newDownloads;
+};
+
+// Complete or cancel a download
+const finishDownload = (id: string) => {
+  const newDownloads = { ...downloadingItems.value };
+  delete newDownloads[id];
+  downloadingItems.value = newDownloads;
+};
+
+// Check if an item is currently downloading
+const isDownloading = (id: string): boolean => {
+  return !!downloadingItems.value[id];
+};
+
+// Get the progress of a downloading item
+const getDownloadProgress = (id: string): number => {
+  return downloadingItems.value[id]?.progress || 0;
+};
+
+// Event listeners
+const onSelectedModelChanged = (callback: (modelId: string | null) => void) => {
+  effect(() => {
+    callback(selectedModel.value);
+  });
+};
+
+const onSelectedLoraChanged = (callback: (loraId: string | null) => void) => {
+  effect(() => {
+    callback(selectedLora.value);
+  });
+};
+
+const onDownloadsChanged = (
+  callback: (downloads: Record<string, DownloadInfo>) => void,
+) => {
+  effect(() => {
+    callback(downloadingItems.value);
+  });
+};
+
+// EXPORTS
+export const dispatchers = {
+  setSelectedModel,
+  setSelectedLora,
+  startDownload,
+  updateDownloadProgress,
+  finishDownload,
+  isDownloading,
+  getDownloadProgress,
+};
+
+export const events = {
+  onSelectedModelChanged,
+  onSelectedLoraChanged,
+  onDownloadsChanged,
+};