f16

Cargo.lock

@@ -665,6 +665,7 @@ dependencies = [
  "candle-transformers",
  "cudarc",
  "filesys",
+ "half",
  "hf-hub 0.3.2",
  "image",
  "libloading 0.8.6",
@@ -1816,6 +1817,7 @@ dependencies = [
  "candle-flash-attn",
  "candle-nn",
  "fancy-regex",
+ "half",
  "intel-mkl-src",
  "num-traits",
  "rand 0.8.5",

app_config.yaml

@@ -1,7 +1,8 @@
-image_width: 512
-image_height: 512
+image_width: 1024
+image_height: 1024
 #scheduler_steps: 15
 # NB(bt): It looks like we can get down to 2 sampling steps, perhaps ?
-scheduler_steps: 2
+scheduler_steps: 8
 scheduler_samples: 1
-cfg_scale: 1.1
+cfg_scale: 15.0
+seed: 42

crates/desktop/src-tauri/Cargo.toml

@@ -45,10 +45,11 @@ serde_json = "1.0"
 tauri = { version = "2.2.4", features = [] }
 tauri-plugin-log = "2"
 tokenizers = "0.15.2"
+half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
 
 [features]
 default = ["cuda"]
-cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
+cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda", "candle-transformers/flash-attn", "candle-core/cudnn"]
 
 [dev-dependencies]
 libloading = "0.8.6"

crates/desktop/src-tauri/src/endpoints/image_endpoint.rs

@@ -1,6 +1,6 @@
 use crate::ml::model_cache::ModelCache;
 use crate::ml::prompt_cache::PromptCache;
-use crate::ml::stable_diffusion::stable_diffusion_pipeline::{stable_diffusion_pipeline, Args};
+use crate::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
 use crate::state::app_config::AppConfig;
 use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
 use bytes::BytesMut;
@@ -13,10 +13,12 @@ use tauri::{AppHandle, State};
 use crate::ml::downloads::download_all_models::download_all_models;
 
 const PROMPT_FILENAME : &str = "prompt.txt";
-const PROMPT: &str = "A beautiful landscape with mountains and a lake";
+const PROMPT: &str = "green goblin, blood dripping from lips, detailed, photorealistic, 8k";
 
 const NEGATIVE_PROMPT: &str = "bad quality, bad faces, poor quality, blurry faces, watermark";
 
+const RANDOM_SEED: u32 = 42;
+
 // TODO: Mark async https://github.com/tauri-apps/tauri/discussions/7737
 
 
@@ -77,9 +79,10 @@ async fn do_infer_image(
     i2i_strength: strength,
     cfg_scale: config.cfg_scale,
     app: &app,
+    use_flash_attn: true,
   };
 
-  match stable_diffusion_pipeline(args) {
+  match lcm_pipeline(args) {
     Ok(image) => {
       let mut bytes = Vec::with_capacity(1024*1024);
       

crates/desktop/src-tauri/src/ml/model_file.rs

@@ -4,80 +4,177 @@ extern crate accelerate_src;
 extern crate intel_mkl_src;
 
 use candle_transformers::models::stable_diffusion;
+use anyhow::{Error as E, Result};
+use hf_hub::api::sync::Api;
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum StableDiffusionVersion {
+    V1_5,
+    V1_5Inpaint,
+    V2_1,
+    V2Inpaint,
+    Xl,
+    XlInpaint,
+    Turbo,
+}
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum ModelFile {
-  Tokenizer,
-  Tokenizer2,
-  Clip,
-  Clip2,
-  Unet,
-  Vae,
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum StableDiffusionVersion {
-  V1_5,
-  V1_5Inpaint,
-  V2_1,
-  V2Inpaint,
-  Xl,
-  XlInpaint,
-  Turbo,
+    Tokenizer,
+    Tokenizer2,
+    Clip,
+    Clip2,
+    Unet,
+    UnetLcm,
+    Vae,
 }
 
 impl StableDiffusionVersion {
-  pub 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",
+    pub 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 => "Lykon/dreamshaper-7",
+            Self::V1_5Inpaint => "stable-diffusion-v1-5/stable-diffusion-inpainting",
+            Self::Turbo => "Lykon/dreamshaper-7",
+        }
+    }
+
+    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 {
-  pub fn get(
-    &self,
-    filename: Option<String>,
-    version: StableDiffusionVersion,
-    use_f16: bool,
-  ) -> anyhow::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 => {
-                "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(), "text_encoder/model.safetensors"),
-          Self::Clip2 => (version.repo(), "text_encoder_2/model.safetensors"),
-          Self::Unet => (version.repo(), "unet/diffusion_pytorch_model.safetensors"),
-          Self::Vae => (version.repo(), "vae/diffusion_pytorch_model.safetensors"),
-        };
-        let filename = Api::new()?.model(repo.to_string()).get(path)?;
-        Ok(filename)
-      }
+    pub fn get(
+        &self,
+        filename: Option<String>,
+        version: StableDiffusionVersion,
+        use_f16: bool,
+    ) -> Result<std::path::PathBuf> {
+        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)
+            }
+        }
     }
-  }
 }
-

crates/desktop/src-tauri/src/ml/models/unet_model.rs

@@ -8,28 +8,23 @@ pub struct UNetModel {
 }
 
 impl UNetModel {
-  pub fn new<P: AsRef<Path>>(
-    sd_config: &StableDiffusionConfig,
-    unet_weights_file: P,
-    device: &Device,
-    dtype: DType,
-  ) -> anyhow::Result<Self> {
+  pub fn new<P: AsRef<Path>>(sd_config: &StableDiffusionConfig, unet_weights_file: P, device: &Device, dtype: DType) -> anyhow::Result<Self> {
     println!("building unet model... (1)");
-    
-    let model = sd_config.build_unet(
-      unet_weights_file, 
-      &device, 
-      4, 
-      false, 
-      dtype
-    )?;
-    
-    Ok(Self {
-      model,
-    })
+
+    let model = sd_config.build_unet(unet_weights_file, &device, 4, true, dtype)?;
+
+    Ok(Self { model })
   }
-  
+
+  pub fn new_with_inner(model: UNet2DConditionModel) -> Self {
+    Self { model }
+  }
+
   pub fn inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result<Tensor> {
     Ok(self.model.forward(&latent_model_input, timestep, &text_embeddings)?)
   }
+
+  pub fn forward_with_guidance(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor, guidance_scale_embedding: &Tensor) -> anyhow::Result<Tensor> {
+    Ok(self.model.forward_with_guidance(&latent_model_input, timestep, &text_embeddings, Some(guidance_scale_embedding))?)
+  }
 }

crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs

@@ -27,7 +27,7 @@ pub fn infer_clip_text_embeddings(
   dtype: DType,
   use_guide_scale: bool,
 ) -> anyhow::Result<Tensor> {
-  info!("Preparing text embeddings...");
+  info!("Preparing text embeddings... for {:?}", sd_version);
 
   let which = match sd_version {
     StableDiffusionVersion::Xl
@@ -43,7 +43,7 @@ pub fn infer_clip_text_embeddings(
         &prompt,
         &uncond_prompt,
         None, // tokenizer
-        None, // clip_weights
+        clip_weights.clone(),
         None, // clip2_weights
         sd_version,
         &sd_config,
@@ -76,6 +76,7 @@ fn do_infer_clip_text_embeddings(
     use_guide_scale: bool,
     first: bool,
   ) -> anyhow::Result<Tensor> {
+    info!("do_infer_clip_text_embeddings called with args {:?}", (prompt, uncond_prompt, tokenizer.clone().unwrap_or_else(|| "None".to_string()), clip_weights.clone().unwrap_or_else(|| "None".to_string()), clip2_weights.clone().unwrap_or_else(|| "None".to_string()), sd_version, sd_config, use_f16, device, dtype, use_guide_scale, first));
   
   let tokenizer_file = if first {
     info!("ModelFile::Tokenizer");
@@ -85,9 +86,9 @@ fn do_infer_clip_text_embeddings(
     ModelFile::Tokenizer2
   };
   
+  let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
   info!("Loading clip from download");
-  //let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
-  let tokenizer = PathBuf::from(ModelRegistry::Clip.get_filename());
+
   
   info!("Tokenizer path: {:?}", tokenizer);
   

crates/desktop/src-tauri/src/ml/stable_diffusion/lcm_pipeline.rs

@@ -0,0 +1,317 @@
+use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
+
+use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
+use crate::events::notification_event::{ModelType, NotificationEvent};
+use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
+use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
+use crate::ml::model_cache::ModelCache;
+use crate::ml::model_registry::ModelRegistry;
+use crate::ml::models::unet_model::UNetModel;
+use crate::ml::prompt_cache::PromptCache;
+use candle_nn::VarBuilder;
+use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
+use candle_transformers::models::stable_diffusion::unet_2d::BlockConfig;
+use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModelConfig;
+use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
+use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings;
+use crate::state::app_config::AppConfig;
+use anyhow::{anyhow, Error as E, Result};
+use candle_core::{DType, IndexOp, Tensor, D};
+use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
+use candle_transformers::models::stable_diffusion::lcm::LCMScheduler;
+use image::DynamicImage;
+use log::info;
+use rand::Rng;
+use tauri::{AppHandle, Emitter};
+
+pub struct Args<'a> {
+  pub image: &'a DynamicImage,
+  pub prompt: String,
+  pub uncond_prompt: String,
+  pub cfg_scale: Option<f64>,
+  pub i2i_strength: Option<u8>,
+  pub configs: &'a AppConfig,
+  pub model_cache: &'a ModelCache,
+  pub prompt_cache: &'a PromptCache,
+  pub app: &'a AppHandle,
+  pub use_flash_attn: bool,
+}
+
+pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
+  let img2img_strength = args.i2i_strength.unwrap_or(75) as f64 / 100.0;
+  let Args { prompt, uncond_prompt, cfg_scale, i2i_strength, configs, model_cache, prompt_cache, app, image, use_flash_attn } = args;
+
+  let use_f16 = true;
+
+  println!("Starting image generation with the following configuration:");
+  println!("  Model: {:?}", configs.sd_version);
+  println!("  Prompt: {}", prompt);
+  println!("  Steps: {}", configs.scheduler_steps);
+  println!("  Device: {:?}", configs.device);
+
+  println!("Model dimensions: {}x{}", configs.sd_config.width, configs.sd_config.height);
+
+  // Use LCM Scheduler instead of Euler Ancestral for better speed and quality
+  let mut scheduler = LCMScheduler::new(configs.scheduler_steps, img2img_strength, candle_transformers::models::stable_diffusion::lcm::LCMSchedulerConfig::default())?;
+
+  let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
+  configs.device.set_seed(seed)?;
+
+  info!("Using seed: {}", seed);
+
+  let guidance_scale = match cfg_scale {
+    Some(guidance_scale) => guidance_scale,
+    None => match configs.sd_version {
+      StableDiffusionVersion::V1_5 | StableDiffusionVersion::V2_1 | StableDiffusionVersion::Xl => 7.5,
+      StableDiffusionVersion::Turbo => 0.,
+      _ => 0., // NB: Not sure what the other model families should use, so sticking with "0"
+    },
+  };
+
+  // let use_guide_scale = guidance_scale > 1.0;
+  let use_guide_scale = false;
+
+  info!("Using guide scale: {}", use_guide_scale);
+
+  info!("Checking if prompt is cached");
+  let maybe_cached = prompt_cache.get_copy(&prompt)?;
+
+  let mut text_embeddings = if let Some(tensor) = maybe_cached {
+    tensor
+  } else {
+    println!("Loading clip weights for {:?} with model {:?}", configs.sd_version, configs.sd_version.repo());
+    let api = configs.hf_api.clone();
+    let clip_path = api.model(configs.sd_version.repo().to_string()).get(if use_f16 { "text_encoder/model.fp16.safetensors" } else { "text_encoder/model.safetensors" })?;
+    let clip_weights = Some(clip_path.to_string_lossy().to_string());
+    info!("Prompt is NOT cached! Calculating embedding...");
+    let tensor = infer_clip_text_embeddings(
+      &prompt,
+      &uncond_prompt,
+      None, // tokenizer
+      clip_weights,
+      None, // clip2_weights
+      configs.sd_version,
+      &configs.sd_config,
+      use_f16, // use_f16
+      &configs.device,
+      configs.dtype,
+      use_guide_scale,
+    )?;
+    prompt_cache.store_copy(&prompt, &tensor)?;
+    tensor
+  };
+
+  println!("Text embeddings shape: {:?}", text_embeddings.shape());
+
+  println!("Loading input image into tensor...");
+
+  let input_image = dynamic_image_to_tensor(image, &configs.device, configs.dtype)?;
+
+  println!("Reference image shape: {:?}", input_image.shape());
+
+  let vae_scale = get_vae_scale(configs.sd_version);
+
+  let maybe_vae = model_cache.get_vae()?;
+
+  let vae = match maybe_vae {
+    Some(vae) => vae,
+    None => {
+      info!("No vae found in cache; loading...");
+
+      let repo = configs.sd_version.repo();
+
+      println!("Building VAE model from : {:?} ... (3)", repo);
+
+      let vae_file = configs.hf_api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?;
+
+      println!("Building VAE model from file {:?}...", &vae_file);
+
+      let mut notify_download_complete = false;
+      //if !vae_file.exists() {
+      if true {
+        notify_download_complete = true;
+        app.emit("notification", NotificationEvent::ModelDownloadStarted { model_name: repo, model_type: ModelType::Vae })?;
+      }
+
+      let vae = configs.sd_config.build_vae(vae_file, &configs.device, configs.dtype)?;
+
+      let vae = Arc::new(vae);
+
+      model_cache.set_vae(vae.clone())?;
+
+      if notify_download_complete {
+        app.emit("notification", NotificationEvent::ModelDownloadComplete { model_name: repo, model_type: ModelType::Vae })?;
+      }
+
+      vae
+    },
+  };
+
+  let maybe_unet = model_cache.get_unet()?;
+
+  let unet = match maybe_unet {
+    Some(unet) => unet,
+    None => {
+      info!("No unet found in cache; loading...");
+
+      let mut notify_download_complete = false;
+      //if !unet_file.exists() {
+      // if true {
+      //     notify_download_complete = true;
+      //     app.emit("notification", NotificationEvent::ModelDownloadStarted {
+      //         model_name: "sdxl-turbo-unet",
+      //         model_type: ModelType::Unet,
+      //     })?;
+      // }
+      let unet_weights = {
+        let lcm_model_repo = "SimianLuo/LCM_Dreamshaper_v7"; // LCM UNet
+        let unet_path = configs.hf_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())
+      };
+
+      let unet_weights = ModelFile::UnetLcm.get(unet_weights, configs.sd_version, use_f16)?;
+      let in_channels = match configs.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
+        time_cond_proj_dim: Some(256), // LCM models use 256-dimensional guidance embeddings
+        // 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,
+        time_embed_dim: Some(1280),
+        // time_embed_dim: Some(256),
+      };
+
+      // Load the model directly from the safetensors file
+      let vs_unet = unsafe { VarBuilder::from_mmaped_safetensors(&[unet_weights], configs.dtype, &configs.device)? };
+
+      let unet = UNet2DConditionModel::new(vs_unet, in_channels, 4, use_flash_attn, unet_config)?;
+
+      let unet = UNetModel::new_with_inner(unet);
+
+      // let unet_file = ModelRegistry::SdxlTurboUnet.get_filename();
+
+      // let unet = UNetModel::new(&configs.sd_config, unet_file, &configs.device, configs.dtype)
+      //   .map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
+
+      let unet = Arc::new(unet);
+
+      model_cache.set_unet(unet.clone())?;
+
+      if notify_download_complete {
+        app.emit("notification", NotificationEvent::ModelDownloadComplete { model_name: "sdxl-turbo-unet", model_type: ModelType::Unet })?;
+      }
+
+      unet
+    },
+  };
+
+  let init_latent_dist: DiagonalGaussianDistribution = vae.encode(&input_image)?;
+
+  // Parse and calculate the img2img strength
+  let img2img_strength = match i2i_strength {
+    None => 0.75f64,
+    Some(strength) => (strength as f64) / 100.0f64,
+  };
+
+  // Create LCM guidance scale embeddings
+  let embedding_dim = 1280;
+  let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &configs.device, configs.dtype)?;
+
+  // Generate latents from input image using the LCM approach
+  println!("Generating latents from input image...");
+  let init_latents = (init_latent_dist.sample()? * vae_scale)?;
+  println!("Initial latents shape: {:?}", init_latents.shape());
+
+  // Calculate timesteps for LCM with proper img2img handling
+  let timesteps = LCMScheduler::get_timesteps_for_steps(configs.scheduler_steps, img2img_strength);
+  let t_start = timesteps[0]; // First denoising step
+
+  // Add noise at the appropriate timestep for img2img
+  println!("Adding noise to latents for t_start={}", t_start);
+  let noise = init_latents.randn_like(0f64, 1f64)?;
+  let latents = scheduler.add_noise(&init_latents, noise, t_start)?;
+
+  println!("Latents initialized successfully");
+  println!("Initial noise shape: {:?}", latents.shape());
+
+  println!("Starting diffusion process...");
+
+  let mut latents = latents;
+  let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();
+
+  for (timestep_index, &timestep) in timesteps.iter().enumerate() {
+    println!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep);
+
+    let latent_model_input = if use_guide_scale { Tensor::cat(&[&latents, &latents], 0)? } else { latents.clone() };
+
+    let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep);
+
+    // Use the guidance scale embedding in the UNet inference with the proper method name
+    let mut noise_pred = match unet.forward_with_guidance(&latent_model_input, timestep as f64, &text_embeddings, &guidance_scale_embedding) {
+      Ok(pred) => pred,
+      Err(e) => {
+        println!("UNet inference failed with error: {}", e);
+        return Err(anyhow::anyhow!("UNet inference failed: {}", e));
+      },
+    };
+
+    if use_guide_scale {
+      let chunks = noise_pred.chunk(2, 0)?;
+      let (noise_pred_uncond, noise_pred_text) = (&chunks[0], &chunks[1]);
+      noise_pred = (noise_pred_uncond + (guidance_scale * (noise_pred_text - noise_pred_uncond)?)?)?;
+    }
+
+    // Apply the LCM scheduler step
+    latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, configs.scheduler_steps)?;
+  }
+
+  println!("Diffusion process completed, decoding image...");
+  let image = vae.decode(&(latents / vae_scale)?)?;
+
+  println!("VAE decode completed");
+
+  println!("Scaling image back");
+  let image = ((image / 2.)? + 0.5)?;
+
+  println!("Normalized image values");
+  let image = (image.clamp(0f32, 1.)? * 255.)?;
+
+  println!("Convert to int8");
+  let image = image.to_dtype(DType::U8)?;
+
+  println!("Converted to 8-bit format");
+
+  let image = tensor_to_image_buffer(&image.i(0)?)?;
+
+  Ok(image)
+}

crates/desktop/src-tauri/src/ml/stable_diffusion/mod.rs

@@ -1,3 +1,4 @@
 pub mod get_vae_scale;
 pub mod infer_clip_text_embeddings;
-pub mod stable_diffusion_pipeline;
+pub mod stable_diffusion_pipeline;
+pub mod lcm_pipeline;

crates/desktop/src-tauri/src/state/app_config.rs

@@ -8,11 +8,11 @@ use log::info;
 
 const CONFIG_FILENAME : &str = "app_config.yaml";
 
-const DEFAULT_SD_IMAGE_WIDTH: usize = 512;
-const DEFAULT_SD_IMAGE_HEIGHT: usize = 512;
+const DEFAULT_SD_IMAGE_WIDTH: usize = 1024;
+const DEFAULT_SD_IMAGE_HEIGHT: usize = 1024;
 
 const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32;
-const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::Turbo;
+const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::V1_5;
 
 /// A central place to configure app-wide details.
 pub struct AppConfig {
@@ -52,14 +52,21 @@ impl AppConfig {
           info!("CUDA not available ({}), falling back to CPU", e);
           Device::Cpu
         });
+
     
     let sd_config = match SD_VERSION {
       StableDiffusionVersion::Turbo => {
         StableDiffusionConfig::sdxl_turbo(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
       }
-      _ => {
+      StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => {
+        StableDiffusionConfig::v1_5(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
+      StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint => {
         StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
       }
+      StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint => {
+        StableDiffusionConfig::sdxl(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
     };
     
     let hf_api = Api::new()?;
@@ -71,7 +78,7 @@ impl AppConfig {
       sd_config,
       image_height: yaml_configs.image_height.unwrap_or(DEFAULT_SD_IMAGE_HEIGHT),
       image_width: yaml_configs.image_width.unwrap_or(DEFAULT_SD_IMAGE_WIDTH),
-      scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(1),
+      scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(4),
       scheduler_samples: yaml_configs.scheduler_samples.unwrap_or(15),
       seed: yaml_configs.seed,
       cfg_scale: yaml_configs.cfg_scale,

crates/vendor/candle/candle-flash-attn/build.rs

@@ -86,6 +86,8 @@ fn main() -> Result<()> {
         .arg("--expt-relaxed-constexpr")
         .arg("--expt-extended-lambda")
         .arg("--use_fast_math")
+        .arg("--compiler-options")
+        .arg("-fPIC")
         .arg("--verbose");
 
     if let Ok(target) = std::env::var("TARGET") {

crates/vendor/candle/candle-nn/Cargo.toml

@@ -32,3 +32,4 @@ metal = ["candle/metal", "dep:candle-metal-kernels", "dep:metal"]
 [[bench]]
 name = "bench_main"
 harness = false
+

crates/vendor/candle/candle-transformers/Cargo.toml

@@ -17,6 +17,7 @@ serde = { version = "1.0.171", features = ["derive"] }
 serde_json = "1.0.99"
 serde_plain = "1.0.2"
 tracing = "0.1.37"
+half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
 
 [features]
 default = []

crates/vendor/candle/candle-transformers/src/models/stable_diffusion/mod.rs

@@ -84,6 +84,8 @@ impl StableDiffusionConfig {
         };
         // https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/unet/config.json
         let unet = unet_2d::UNet2DConditionModelConfig {
+            time_embed_dim: None,
+            time_cond_proj_dim: None,
             blocks: vec![
                 bc(320, Some(1), 8),
                 bc(640, Some(1), 8),
@@ -153,6 +155,8 @@ impl StableDiffusionConfig {
         };
         // https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/unet/config.json
         let unet = unet_2d::UNet2DConditionModelConfig {
+            time_cond_proj_dim: None,
+            time_embed_dim: None,
             blocks: vec![
                 bc(320, Some(1), 5),
                 bc(640, Some(1), 10),
@@ -237,6 +241,8 @@ impl StableDiffusionConfig {
         };
         // https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/unet/config.json
         let unet = unet_2d::UNet2DConditionModelConfig {
+            time_cond_proj_dim: None,
+            time_embed_dim: None,
             blocks: vec![
                 bc(320, None, 5),
                 bc(640, Some(2), 10),
@@ -306,6 +312,8 @@ impl StableDiffusionConfig {
         };
         // https://huggingface.co/stabilityai/sdxl-turbo/blob/main/unet/config.json
         let unet = unet_2d::UNet2DConditionModelConfig {
+            time_cond_proj_dim: None,
+            time_embed_dim: None,
             blocks: vec![
                 bc(320, None, 5),
                 bc(640, Some(2), 10),
@@ -405,6 +413,8 @@ impl StableDiffusionConfig {
         };
         // https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/unet/config.json
         let unet = unet_2d::UNet2DConditionModelConfig {
+            time_cond_proj_dim: None,
+            time_embed_dim: None,
             blocks: vec![
                 bc(320, None, 5),
                 bc(640, Some(2), 10),

crates/vendor/candle/candle-transformers/src/models/stable_diffusion/unet_2d.rs

@@ -10,6 +10,10 @@ use candle_nn as nn;
 use candle_nn::Module;
 use std::collections::HashMap;
 
+// fn save_tensor_as_numpy(tensor: &Tensor, filename: &str) -> Result<()> {
+//     // Use Candle's built-in NPY serialization
+//     tensor.write_npy(filename).map_err(|e| anyhow::anyhow!("{}", e))
+// }
 #[derive(Debug, Clone, Copy)]
 pub struct BlockConfig {
   pub out_channels: usize,
@@ -33,11 +37,13 @@ pub struct UNet2DConditionModelConfig {
   pub cross_attention_dim: usize,
   pub sliced_attention_size: Option<usize>,
   pub use_linear_projection: bool,
+  pub time_cond_proj_dim: Option<usize>,
+  pub time_embed_dim: Option<usize>,
 }
 
 impl Default for UNet2DConditionModelConfig {
   fn default() -> Self {
-    Self { center_input_sample: false, flip_sin_to_cos: true, freq_shift: 0., blocks: vec![BlockConfig { out_channels: 320, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 640, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: None, attention_head_dim: 8 }], layers_per_block: 2, downsample_padding: 1, mid_block_scale_factor: 1., norm_num_groups: 32, norm_eps: 1e-5, cross_attention_dim: 1280, sliced_attention_size: None, use_linear_projection: false }
+    Self { center_input_sample: false, flip_sin_to_cos: true, freq_shift: 0., blocks: vec![BlockConfig { out_channels: 320, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 640, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: None, attention_head_dim: 8 }], layers_per_block: 2, downsample_padding: 1, mid_block_scale_factor: 1., norm_num_groups: 32, norm_eps: 1e-5, cross_attention_dim: 1280, sliced_attention_size: None, use_linear_projection: false, time_cond_proj_dim: Some(256), time_embed_dim: None }
   }
 }
 
@@ -73,8 +79,16 @@ impl UNet2DConditionModel {
     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 time_embed_dim = config.time_embed_dim.unwrap_or(b_channels * 4);
     let conv_cfg = nn::Conv2dConfig { padding: 1, ..Default::default() };
+
+    // Log important configuration details
+    println!("Creating UNet2DConditionModel with {} blocks", n_blocks);
+    if let Some(proj_dim) = config.time_cond_proj_dim {
+      println!("UNet configured with time_cond_proj_dim: {}", proj_dim);
+      println!("This means the model expects guidance embeddings of dimension {}", proj_dim);
+    }
+
     let conv_in = conv2d(in_channels, b_channels, 3, conv_cfg, vs.pp("conv_in"))?;
 
     let time_proj = Timesteps::new(b_channels, config.flip_sin_to_cos, config.freq_shift);
@@ -241,18 +255,138 @@ impl UNet2DConditionModel {
 
     // 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 dimensions of both embeddings
+      let (batch_size, time_emb_dim) = time_emb.dims2()?;
+      let (guidance_batch, guidance_emb_dim) = guidance_emb.dims2()?;
 
-      // Get the batch size and embedding dimension
-      let (batch_size, emb_dim) = time_emb.dims2()?;
+      println!("==== Guidance Embedding Integration ====");
+      println!("Time embedding dims: [{}, {}]", batch_size, time_emb_dim);
+      println!("Guidance embedding dims: [{}, {}]", guidance_batch, guidance_emb_dim);
 
-      // Reshape guidance embedding to match time embedding dimensions
-      let guidance_emb = guidance_emb.reshape(&[1, emb_dim])?.expand(&[batch_size, emb_dim])?;
+      // Use time_cond_proj_dim from config to determine expected guidance embedding dimension
+      let expected_guidance_dim = self.config.time_cond_proj_dim.unwrap_or(time_emb_dim);
+      println!("Expected guidance dimension from config: {}", expected_guidance_dim);
+
+      // Check if dimensions match the expected projection dimension
+      let processed_guidance_emb = if guidance_emb_dim != time_emb_dim {
+        // Convert guidance_emb to f32 for calculations
+        let guidance_emb_f32 = guidance_emb.to_dtype(DType::F32)?;
+        println!("Converting guidance embedding to F32 for processing");
+
+        // If the guidance embedding dimension matches the expected projection dimension,
+        // we need to project it to match the time embedding dimension
+        if guidance_emb_dim == expected_guidance_dim {
+          println!("Projecting guidance embedding from {} to {}", guidance_emb_dim, time_emb_dim);
+
+          // Use linear interpolation
+          let src_positions: Vec<f32> = (0..guidance_emb_dim).map(|i| i as f32 / (guidance_emb_dim - 1) as f32).collect();
+          let dst_positions: Vec<f32> = (0..time_emb_dim).map(|i| i as f32 / (time_emb_dim - 1) as f32).collect();
+
+          // Get guidance embedding data (now f32)
+          let guidance_data = guidance_emb_f32.reshape((guidance_emb_dim,))?.to_vec1::<f32>()?;
+          println!("Reshaping guidance data to vector of length {}", guidance_data.len());
+
+          // Perform linear interpolation
+          let mut interpolated_data = Vec::with_capacity(time_emb_dim);
+          for &t in &dst_positions {
+            // Find the two source positions that surround the target position
+            let mut i = 0;
+            while i < src_positions.len() - 1 && src_positions[i + 1] < t {
+              i += 1;
+            }
+
+            if i == src_positions.len() - 1 {
+              // At or beyond the last position
+              interpolated_data.push(guidance_data[i]);
+            } else if t <= src_positions[0] {
+              // At or before the first position
+              interpolated_data.push(guidance_data[0]);
+            } else {
+              // Linear interpolation between two points
+              let t_norm = (t - src_positions[i]) / (src_positions[i + 1] - src_positions[i]);
+              let value = guidance_data[i] * (1.0 - t_norm) + guidance_data[i + 1] * t_norm;
+              interpolated_data.push(value);
+            }
+          }
+
+          println!("Created interpolated data of length {}", interpolated_data.len());
+
+          // Create tensor and convert to time_emb's dtype
+          let tensor = Tensor::from_vec(interpolated_data, (1, time_emb_dim), device)?;
+          let result = tensor.to_dtype(time_emb.dtype())?;
+          println!("Final projected guidance embedding shape: {:?}", result.shape());
+          result
+        } else {
+          // Unusual case - dimensions don't match expected pattern
+          println!("Warning: Guidance embedding dimension ({}) doesn't match expected projection dimension ({})", guidance_emb_dim, expected_guidance_dim);
+
+          // Fallback to generic interpolation
+          let guidance_data = guidance_emb_f32.reshape((guidance_emb_dim,))?.to_vec1::<f32>()?;
+
+          let src_positions: Vec<f32> = (0..guidance_emb_dim).map(|i| i as f32 / (guidance_emb_dim - 1) as f32).collect();
+          let dst_positions: Vec<f32> = (0..time_emb_dim).map(|i| i as f32 / (time_emb_dim - 1) as f32).collect();
+
+          let mut interpolated_data = Vec::with_capacity(time_emb_dim);
+          for &t in &dst_positions {
+            let mut i = 0;
+            while i < src_positions.len() - 1 && src_positions[i + 1] < t {
+              i += 1;
+            }
+
+            if i == src_positions.len() - 1 {
+              interpolated_data.push(guidance_data[i]);
+            } else if t <= src_positions[0] {
+              interpolated_data.push(guidance_data[0]);
+            } else {
+              let t_norm = (t - src_positions[i]) / (src_positions[i + 1] - src_positions[i]);
+              let value = guidance_data[i] * (1.0 - t_norm) + guidance_data[i + 1] * t_norm;
+              interpolated_data.push(value);
+            }
+          }
+
+          println!("Created interpolated data of length {}", interpolated_data.len());
+          let tensor = Tensor::from_vec(interpolated_data, (1, time_emb_dim), device)?;
+          let result = tensor.to_dtype(time_emb.dtype())?;
+          println!("Final fallback guidance embedding shape: {:?}", result.shape());
+          result
+        }
+      } else {
+        // Dimensions match, no processing needed
+        println!("Guidance dimension already matches time embedding dimension, no interpolation needed");
+        let result = guidance_emb.reshape((1, time_emb_dim))?.to_dtype(time_emb.dtype())?;
+        println!("Reshaped guidance embedding shape: {:?}", result.shape());
+        result
+      };
+
+      // Expand to match batch size if needed
+      let expanded_guidance = processed_guidance_emb.broadcast_as(time_emb.shape())?;
+      println!("Expanded guidance embedding shape: {:?}", expanded_guidance.shape());
+
+      // Save the processed and expanded guidance embeddings to disk if requested
+      if let Ok(save_dir) = std::env::var("CANDLE_SAVE_TENSORS_DIR") {
+        // Save pre-expansion (original projection)
+        let proc_path = format!("{}/unet_processed_guidance_embedding.npy", save_dir);
+        println!("Saving processed guidance embedding to {}", proc_path);
+        if let Err(e) = processed_guidance_emb.write_npy(&proc_path) {
+          println!("Error saving processed guidance embedding: {}", e);
+        } else {
+          println!("Successfully saved processed guidance embedding to {}", proc_path);
+        }
+
+        // Save post-expansion (final tensor used for addition to time_emb)
+        let exp_path = format!("{}/unet_expanded_guidance_embedding.npy", save_dir);
+        println!("Saving expanded guidance embedding to {}", exp_path);
+        if let Err(e) = expanded_guidance.write_npy(&exp_path) {
+          println!("Error saving expanded guidance embedding: {}", e);
+        } else {
+          println!("Successfully saved expanded guidance embedding to {}", exp_path);
+        }
+      }
 
       // Add to time embeddings
-      time_emb = (time_emb + guidance_emb)?;
+      time_emb = (time_emb + expanded_guidance)?;
+      println!("Final time embedding with guidance shape: {:?}", time_emb.shape());
+      println!("==== End Guidance Integration ====");
     }
 
     // 2. pre-process