don't download sdxl anymore

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

@@ -9,6 +9,9 @@ pub enum ModelType {
   SdxlTurboVae,
   SdxlTurboClipEncoder,
   SdxlTurboClipEncoder2,
+  SimianLuoLcmDreamshaperV7Unet,
+  LykonDreamshaper7Vae,
+  LykonDreamshaper7TextEncoderFp16,
 }
 
 impl ModelType {
@@ -19,6 +22,9 @@ impl ModelType {
       Self::SdxlTurboVae => "SDXL Turbo VAE",
       Self::SdxlTurboClipEncoder => "Clip Encoder #1",
       Self::SdxlTurboClipEncoder2 => "Clip Encoder #2",
+      Self::SimianLuoLcmDreamshaperV7Unet => "SimianLuo LCM Dreamshaper v7 UNET",
+      Self::LykonDreamshaper7Vae => "Lykon Dreamshaper 7 VAE",
+      Self::LykonDreamshaper7TextEncoderFp16 => "Lykon Dreamshaper Text Encoder (fp16)",
     }
   }
 
@@ -29,6 +35,9 @@ impl ModelType {
       Self::SdxlTurboVae => NotificationModelType::Vae,
       Self::SdxlTurboClipEncoder => NotificationModelType::ClipEncoder,
       Self::SdxlTurboClipEncoder2 => NotificationModelType::ClipEncoder,
+      Self::SimianLuoLcmDreamshaperV7Unet => NotificationModelType::Unet,
+      Self::LykonDreamshaper7Vae => NotificationModelType::Vae,
+      Self::LykonDreamshaper7TextEncoderFp16 => NotificationModelType::ClipEncoder,
     }
   }
 
@@ -39,6 +48,9 @@ impl ModelType {
       | Self::SdxlTurboVae
       | Self::SdxlTurboClipEncoder
       | Self::SdxlTurboClipEncoder2 => Some("stabilityai/sdxl-turbo"),
+      Self::SimianLuoLcmDreamshaperV7Unet => Some("SimianLuo/LCM_Dreamshaper_v7"),
+      Self::LykonDreamshaper7Vae
+      | Self::LykonDreamshaper7TextEncoderFp16 => Some("Lykon/dreamshaper-7"),
     }
   }
 
@@ -49,6 +61,9 @@ impl ModelType {
       Self::SdxlTurboVae => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/diffusion_pytorch_model.vae.safetensors",
       Self::SdxlTurboClipEncoder => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder.safetensors",
       Self::SdxlTurboClipEncoder2 => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder_2.safetensors",
+      Self::SimianLuoLcmDreamshaperV7Unet => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/simianluo_lcm_dreamshaper_v7_unet.safetensors",
+      Self::LykonDreamshaper7Vae => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_vae.safetensors",
+      Self::LykonDreamshaper7TextEncoderFp16 => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_text_encoder.fp16.safetensors",
     }
   }
 
@@ -59,6 +74,9 @@ impl ModelType {
       Self::SdxlTurboVae => "diffusion_pytorch_model.vae.safetensors",
       Self::SdxlTurboClipEncoder => "clip_text_encoder.safetensors",
       Self::SdxlTurboClipEncoder2 => "clip_text_encoder_2.safetensors",
+      Self::SimianLuoLcmDreamshaperV7Unet => "simianluo_lcm_dreamshaper_v7_unet.safetensors",
+      Self::LykonDreamshaper7Vae => "lykon_dreamshaper_7_vae.safetensors",
+      Self::LykonDreamshaper7TextEncoderFp16 => "lykon_dreamshaper_7_text_encoder.fp16.safetensors",
     }
   }
   

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

@@ -53,20 +53,19 @@ pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer:
 fn do_infer_clip_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, weights_dir: &AppWeightsDir) -> 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");
-    ModelFile::Tokenizer
-  } else {
-    info!("ModelFile::Tokenizer2");
-    ModelFile::Tokenizer2
-  };
+  //let tokenizer_file = if first {
+  //  info!("ModelFile::Tokenizer");
+  //  ModelFile::Tokenizer
+  //} else {
+  //  info!("ModelFile::Tokenizer2");
+  //  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 = weights_dir.model_path(&ModelType::ClipJson);
 
-  info!("Tokenizer path: {:?}", tokenizer);
+  info!("Loading Clip Tokenizer path: {:?}", tokenizer);
 
   let tokenizer = Tokenizer::from_file(tokenizer).map_err(E::msg)?;
 

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

@@ -1,28 +1,30 @@
+use std::path::PathBuf;
 use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
 
-use crate::events::notification_event::{NotificationModelType, NotificationEvent};
-use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
+use crate::events::notification_event::{NotificationEvent, NotificationModelType};
 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_file::{ModelFile, StableDiffusionVersion};
+use crate::ml::model_type::ModelType;
 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 crate::state::app_dir::AppDataRoot;
 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_nn::VarBuilder;
 use candle_transformers::models::stable_diffusion::lcm::LCMScheduler;
+use candle_transformers::models::stable_diffusion::unet_2d::BlockConfig;
+use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
+use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModelConfig;
+use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
 use image::DynamicImage;
 use log::info;
 use rand::Rng;
 use tauri::{AppHandle, Emitter};
-use crate::state::app_dir::AppDataRoot;
 
 pub struct Args<'a> {
   pub image: &'a DynamicImage,
@@ -46,13 +48,13 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
   let img2img_strength = args.i2i_strength.unwrap_or(75) as f64 / 100.0;
   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!("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);
+  //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())?;
@@ -60,7 +62,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
   let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
   configs.device.set_seed(seed)?;
 
-  info!("Using seed: {}", seed);
+  //info!("Using seed: {}", seed);
 
   let guidance_scale = match cfg_scale {
     Some(guidance_scale) => guidance_scale,
@@ -74,18 +76,28 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
   // let use_guide_scale = guidance_scale > 1.0;
   let use_guide_scale = false;
 
-  info!("Using guide scale: {}", use_guide_scale);
+  //info!("Using guide scale: {}", use_guide_scale);
 
-  info!("Checking if prompt is cached");
+  //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!("Loading clip weights...");
+    //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" })?;
+
+    //println!(">>>> CLIP REPO = {:?}", configs.sd_version.repo());
+    //println!(">>>> CLIP PATH = {:?}", clip_path);
+    //println!(">>>> USE_FP16 = {:?}", use_f16);
+    //let clip_weights = Some(clip_path.to_string_lossy().to_string());
+
+    let clip_weights = weights_dir.model_path(&ModelType::LykonDreamshaper7TextEncoderFp16);
+    let clip_weights = Some(clip_weights.to_string_lossy().to_string());
+
     info!("Prompt is NOT cached! Calculating embedding...");
     let tensor = infer_clip_text_embeddings(
       &prompt,
@@ -117,30 +129,27 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
 
   let maybe_vae = model_cache.get_vae()?;
 
-  let repo = configs.sd_version.repo();
+  let vae = match maybe_vae {
+    Some(vae) => vae,
+    None => {
+      //let repo = configs.sd_version.repo();
+      //println!(">>>>> VAE REPO = {:?}", 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 : {:?} ... (3)", repo);
+      let vae_file = weights_dir.model_path(&ModelType::LykonDreamshaper7Vae);
 
-  let vae_file = configs.hf_api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?;
+      println!("Building VAE model from file {:?}...", &vae_file);
 
-  println!("Building VAE model from file {:?}...", &vae_file);
+      let vae = configs.sd_config.build_vae(vae_file, &configs.device, configs.dtype)?;
 
-  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: NotificationModelType::Vae })?;
-  }
+      let vae = Arc::new(vae);
 
-  let vae = configs.sd_config.build_vae(vae_file, &configs.device, configs.dtype)?;
+      model_cache.set_vae(vae.clone())?;
 
-  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: NotificationModelType::Vae })?;
-  }
+      vae
+    }
+  };
 
   let maybe_unet = model_cache.get_unet()?;
 
@@ -149,38 +158,12 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
     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 = weights_dir.model_path(&ModelType::SimianLuoLcmDreamshaperV7Unet);
 
-      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)
@@ -194,7 +177,12 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
         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 }],
+        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,
@@ -213,19 +201,10 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
 
       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: NotificationModelType::Unet })?;
-      }
-
       unet
     },
   };

crates/desktop/src-tauri/src/threads/downloader_thread.rs

@@ -41,10 +41,13 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !
 
   // TODO: Automatic enqueue of known important models + enqueue new models on-demand
   download_queue.push_back(ModelType::ClipJson);
-  download_queue.push_back(ModelType::SdxlTurboUnet);
-  download_queue.push_back(ModelType::SdxlTurboVae);
-  download_queue.push_back(ModelType::SdxlTurboClipEncoder);
-  download_queue.push_back(ModelType::SdxlTurboClipEncoder2);
+  //download_queue.push_back(ModelType::SdxlTurboUnet);
+  //download_queue.push_back(ModelType::SdxlTurboVae);
+  download_queue.push_back(ModelType::SdxlTurboClipEncoder); // TODO(bt): Why is this still needed?
+  //download_queue.push_back(ModelType::SdxlTurboClipEncoder2);
+  download_queue.push_back(ModelType::SimianLuoLcmDreamshaperV7Unet);
+  download_queue.push_back(ModelType::LykonDreamshaper7Vae);
+  download_queue.push_back(ModelType::LykonDreamshaper7TextEncoderFp16);
 
   let mut handles = FuturesUnordered::new();