diff --git a/crates/desktop/src-tauri/src/ml/model_type.rs b/crates/desktop/src-tauri/src/ml/model_type.rs index 69b95acb09..87f14d7c14 100644 --- a/crates/desktop/src-tauri/src/ml/model_type.rs +++ b/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", } } diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs b/crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs index 840c9c3744..d9915b3248 100644 --- a/crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs +++ b/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, clip_weights: Option, clip2_weights: Option, 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 { 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)?; diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/lcm_pipeline.rs b/crates/desktop/src-tauri/src/ml/stable_diffusion/lcm_pipeline.rs index 983612eae8..9bc942b65f 100644 --- a/crates/desktop/src-tauri/src/ml/stable_diffusion/lcm_pipeline.rs +++ b/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 { 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 { 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 { // 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 { 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 { 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 { 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 { 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 }, }; diff --git a/crates/desktop/src-tauri/src/threads/downloader_thread.rs b/crates/desktop/src-tauri/src/threads/downloader_thread.rs index 490588820d..921a061207 100644 --- a/crates/desktop/src-tauri/src/threads/downloader_thread.rs +++ b/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();