From 5a0b2e9d03aed1efe485c8190bd7fdeec1d6d1a2 Mon Sep 17 00:00:00 2001 From: Kasisnu <5044120+kasisnu@users.noreply.github.com> Date: Mon, 10 Mar 2025 13:19:05 +0530 Subject: [PATCH] f16 --- Cargo.lock | 2 + app_config.yaml | 9 +- crates/desktop/src-tauri/Cargo.toml | 5 +- .../src-tauri/src/endpoints/image_endpoint.rs | 5 +- crates/desktop/src-tauri/src/ml/model_file.rs | 227 +++++--- .../src-tauri/src/ml/models/unet_model.rs | 4 + .../infer_clip_text_embeddings.rs | 89 +--- .../src/ml/stable_diffusion/lcm_pipeline.rs | 487 +++++++++--------- .../desktop/src-tauri/src/state/app_config.rs | 15 +- .../candle-examples/examples/lcm/main.rs | 6 +- .../candle/candle-transformers/Cargo.toml | 1 + .../src/models/stable_diffusion/lcm.rs | 10 +- 12 files changed, 473 insertions(+), 387 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3a6806275e..c3ad814075 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -667,6 +667,7 @@ dependencies = [ "dirs 6.0.0", "filesys", "futures 0.3.29", + "half", "hf-hub 0.3.2", "image", "libloading 0.8.6", @@ -1823,6 +1824,7 @@ dependencies = [ "candle-flash-attn", "candle-nn", "fancy-regex", + "half", "intel-mkl-src", "num-traits", "rand 0.8.5", diff --git a/app_config.yaml b/app_config.yaml index 2ddf4fd258..5b3f806ae1 100644 --- a/app_config.yaml +++ b/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: 4 +scheduler_steps: 8 scheduler_samples: 1 -cfg_scale: 1.1 +cfg_scale: 15.0 +seed: 42 \ No newline at end of file diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index 42a7ad4d74..a4596a4d39 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -53,6 +53,7 @@ tempfile = "3.18.0" tokenizers = "0.15.2" tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] } tokio-util = { version = "0.7", features = ["codec"] } +half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] } [features] default = ["cuda"] @@ -61,8 +62,8 @@ cuda = [ "candle-nn/cuda", "candle-transformers/cuda", # NB(bt,2025-03-10): Building without Flash attention on Windows for now to bypass linker issues -# "candle-transformers/flash-attn", -# "candle-core/cudnn" + "candle-transformers/flash-attn", + "candle-core/cudnn" ] [dev-dependencies] diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index ffac1b7003..848a40a515 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs @@ -14,10 +14,12 @@ use tauri::{AppHandle, State}; use crate::state::app_dir::AppDataRoot; 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 @@ -90,6 +92,7 @@ async fn do_infer_image( cfg_scale: config.cfg_scale, app: &app, app_data_root, + use_flash_attn: true, }; match lcm_pipeline(args) { diff --git a/crates/desktop/src-tauri/src/ml/model_file.rs b/crates/desktop/src-tauri/src/ml/model_file.rs index bdfe9b8cc0..4ecba64332 100644 --- a/crates/desktop/src-tauri/src/ml/model_file.rs +++ b/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, - version: StableDiffusionVersion, - use_f16: bool, - ) -> anyhow::Result { - 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, + version: StableDiffusionVersion, + use_f16: bool, + ) -> Result { + 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) + } + } } - } } - diff --git a/crates/desktop/src-tauri/src/ml/models/unet_model.rs b/crates/desktop/src-tauri/src/ml/models/unet_model.rs index 4a8416c02f..cfc1b408fb 100644 --- a/crates/desktop/src-tauri/src/ml/models/unet_model.rs +++ b/crates/desktop/src-tauri/src/ml/models/unet_model.rs @@ -16,6 +16,10 @@ impl UNetModel { 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 { Ok(self.model.forward(&latent_model_input, timestep, &text_embeddings)?) } 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 50965e9d13..840c9c3744 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 @@ -15,26 +15,11 @@ use crate::ml::model_type::ModelType; use crate::state::app_dir::AppWeightsDir; #[allow(clippy::too_many_arguments)] -pub fn 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, - weights_dir: &AppWeightsDir, -) -> anyhow::Result { - info!("Preparing text embeddings..."); +pub fn 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, weights_dir: &AppWeightsDir) -> anyhow::Result { + info!("Preparing text embeddings... for {:?}", sd_version); let which = match sd_version { - StableDiffusionVersion::Xl - | StableDiffusionVersion::XlInpaint - | StableDiffusionVersion::Turbo => vec![true, false], + StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Turbo => vec![true, false], _ => vec![true], }; @@ -45,7 +30,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, @@ -60,27 +45,14 @@ pub fn infer_clip_text_embeddings( .collect::>>()?; let text_embeddings = Tensor::cat(&text_embeddings, D::Minus1)?; - + Ok(text_embeddings) } #[allow(clippy::too_many_arguments)] -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 { - +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 @@ -88,32 +60,25 @@ fn do_infer_clip_text_embeddings( 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); - + 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(); + 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 - ) + 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) @@ -121,11 +86,7 @@ fn do_infer_clip_text_embeddings( 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_file = if first { ModelFile::Clip } else { ModelFile::Clip2 }; let clip_weights = if first { //clip_weights_file.get(clip_weights, sd_version, use_f16)? weights_dir.model_path(&ModelType::SdxlTurboClipEncoder) @@ -133,27 +94,15 @@ fn do_infer_clip_text_embeddings( //clip_weights_file.get(clip2_weights, sd_version, use_f16)? weights_dir.model_path(&ModelType::SdxlTurboClipEncoder2) }; - let clip_config = if first { - &sd_config.clip - } else { - sd_config.clip2.as_ref().unwrap() - }; + 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(); + 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 - ) + 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) 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 f7a5163126..983612eae8 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,13 +1,16 @@ use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; use crate::events::notification_event::{NotificationModelType, NotificationEvent}; +use crate::ml::model_file::{ModelFile, StableDiffusionVersion}; 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::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; @@ -22,272 +25,288 @@ use tauri::{AppHandle, Emitter}; use crate::state::app_dir::AppDataRoot; pub struct Args<'a> { - pub image: &'a DynamicImage, - pub prompt: String, - pub uncond_prompt: String, - pub cfg_scale: Option, - pub i2i_strength: Option, - pub configs: &'a AppConfig, - pub model_cache: &'a ModelCache, - pub prompt_cache: &'a PromptCache, - pub app_data_root: &'a AppDataRoot, - pub app: &'a AppHandle, + pub image: &'a DynamicImage, + pub prompt: String, + pub uncond_prompt: String, + pub cfg_scale: Option, + pub i2i_strength: Option, + pub configs: &'a AppConfig, + pub model_cache: &'a ModelCache, + pub prompt_cache: &'a PromptCache, + pub app_data_root: &'a AppDataRoot, + pub app: &'a AppHandle, + pub use_flash_attn: bool, } pub fn lcm_pipeline(args: Args<'_>) -> Result { - let Args { - prompt, - uncond_prompt, - cfg_scale, - i2i_strength, - configs, - model_cache, - prompt_cache, - app, - image, - app_data_root, - } = args; + let Args { prompt, uncond_prompt, cfg_scale, i2i_strength, configs, model_cache, prompt_cache, app, image, app_data_root, use_flash_attn } = args; - let weights_dir = app_data_root.weights_dir(); - - 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); + let weights_dir = app_data_root.weights_dir(); - println!("Model dimensions: {}x{}", configs.sd_config.width, configs.sd_config.height); + let img2img_strength = args.i2i_strength.unwrap_or(75) as f64 / 100.0; + let use_f16 = true; - // Use LCM Scheduler instead of Euler Ancestral for better speed and quality - let mut scheduler = LCMScheduler::new( - configs.scheduler_steps, - candle_transformers::models::stable_diffusion::lcm::LCMSchedulerConfig::default(), + 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, + weights_dir, )?; + prompt_cache.store_copy(&prompt, &tensor)?; + tensor + }; - let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen()); - configs.device.set_seed(seed)?; + println!("Text embeddings shape: {:?}", text_embeddings.shape()); - info!("Using seed: {}", seed); + println!("Loading input image into tensor..."); - 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 input_image = dynamic_image_to_tensor(image, &configs.device, configs.dtype)?; - let use_guide_scale = guidance_scale > 1.0; + println!("Reference image shape: {:?}", input_image.shape()); - info!("Using guide scale: {}", use_guide_scale); + let vae_scale = get_vae_scale(configs.sd_version); - info!("Checking if prompt is cached"); - let maybe_cached = prompt_cache.get_copy(&prompt)?; + let maybe_vae = model_cache.get_vae()?; - let mut text_embeddings = if let Some(tensor) = maybe_cached { - tensor - } else { - - info!("Prompt is NOT cached! Calculating embedding..."); - let tensor = infer_clip_text_embeddings( - &prompt, - &uncond_prompt, - None, // tokenizer - None, // clip_weights - None, // clip2_weights - configs.sd_version, - &configs.sd_config, - false, // use_f16 - &configs.device, - configs.dtype, - use_guide_scale, - weights_dir, - )?; - prompt_cache.store_copy(&prompt, &tensor)?; - tensor - }; + let repo = configs.sd_version.repo(); - println!("Text embeddings shape: {:?}", text_embeddings.shape()); + println!("Building VAE model from : {:?} ... (3)", repo); - println!("Loading input image into tensor..."); + let vae_file = configs.hf_api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?; - let input_image = dynamic_image_to_tensor( - image, - &configs.device, - configs.dtype)?; + println!("Building VAE model from file {:?}...", &vae_file); - println!("Reference image shape: {:?}", input_image.shape()); + 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_scale = get_vae_scale(configs.sd_version); + let vae = configs.sd_config.build_vae(vae_file, &configs.device, configs.dtype)?; - let maybe_vae = model_cache.get_vae()?; - - let vae = match maybe_vae { - Some(vae) => vae, - None => { - info!("No vae found in cache; loading..."); - - let mut notify_download_complete = false; - //if !vae_file.exists() { - if true { - notify_download_complete = true; - app.emit("notification", NotificationEvent::ModelDownloadStarted { - model_name: "sdxl-turbo-vae", - model_type: NotificationModelType::Vae, - })?; - } + let vae = Arc::new(vae); - let vae_file = weights_dir.model_path(&ModelType::SdxlTurboVae); - - 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: "sdxl-turbo-vae", - model_type: NotificationModelType::Vae, - })?; - } - - vae - } - }; + if notify_download_complete { + app.emit("notification", NotificationEvent::ModelDownloadComplete { model_name: repo, model_type: NotificationModelType::Vae })?; + } - let maybe_unet = model_cache.get_unet()?; + let maybe_unet = model_cache.get_unet()?; - let unet = match maybe_unet { - Some(unet) => unet, - None => { - info!("No unet found in cache; loading..."); + 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: NotificationModelType::Unet, - })?; - } - - let unet_file = weights_dir.model_path(&ModelType::SdxlTurboUnet); - - 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 - } - }; - - 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 = 256; - 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, ×tep) 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)? + 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 { - latents.clone() - }; + "unet/diffusion_pytorch_model.safetensors" + })?; + Some(unet_path.to_string_lossy().to_string()) + }; - let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep); + 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"); - // 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)); - }, - }; + // Ensure we're loading the correct UNet configuration + println!("Building UNet for LCM (SimianLuo/LCM_Dreamshaper_v7)"); - 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)?)?)?; - } + // 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), + }; - // Apply the LCM scheduler step - latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, configs.scheduler_steps)?; + // 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: NotificationModelType::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, ×tep) 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)?)?)?; } - println!("Diffusion process completed, decoding image..."); - let image = vae.decode(&(latents / vae_scale)?)?; - - println!("VAE decode completed"); + // Apply the LCM scheduler step + latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, configs.scheduler_steps)?; + } - println!("Scaling image back"); - let image = ((image / 2.)? + 0.5)?; - - println!("Normalized image values"); - let image = (image.clamp(0f32, 1.)? * 255.)?; + println!("Diffusion process completed, decoding image..."); + let image = vae.decode(&(latents / vae_scale)?)?; - println!("Convert to int8"); - let image = image.to_dtype(DType::U8)?; - - println!("Converted to 8-bit format"); + println!("VAE decode completed"); - let image = tensor_to_image_buffer(&image.i(0)?)?; + println!("Scaling image back"); + let image = ((image / 2.)? + 0.5)?; - Ok(image) + 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) } diff --git a/crates/desktop/src-tauri/src/state/app_config.rs b/crates/desktop/src-tauri/src/state/app_config.rs index 896962daeb..c8ac8de251 100644 --- a/crates/desktop/src-tauri/src/state/app_config.rs +++ b/crates/desktop/src-tauri/src/state/app_config.rs @@ -9,11 +9,11 @@ use hf_hub::api::sync::Api; use log::info; use std::path::PathBuf; -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; const DEFAULT_DATA_DIR_UNIX : &str = "~/artcraft"; @@ -57,14 +57,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()?; diff --git a/crates/vendor/candle/candle-examples/examples/lcm/main.rs b/crates/vendor/candle/candle-examples/examples/lcm/main.rs index 746198475f..55bd3fa55a 100644 --- a/crates/vendor/candle/candle-examples/examples/lcm/main.rs +++ b/crates/vendor/candle/candle-examples/examples/lcm/main.rs @@ -623,6 +623,7 @@ fn run(args: Args) -> Result<()> { let mut scheduler = stable_diffusion::lcm::LCMScheduler::new( n_steps, + strength, stable_diffusion::lcm::LCMSchedulerConfig::default(), )?; @@ -631,7 +632,8 @@ fn run(args: Args) -> Result<()> { 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 use_guide_scale = guidance_scale > 1.0; + let use_guide_scale = false; let which = match sd_version { StableDiffusionVersion::Xl @@ -984,7 +986,7 @@ fn run(args: Args) -> Result<()> { println!("Starting LCM sampling"); let start_time = std::time::Instant::now(); - let embedding_dim = 256; + let embedding_dim = 1280; let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &device, dtype)?; for (timestep_index, ×tep) in timesteps.iter().enumerate() { diff --git a/crates/vendor/candle/candle-transformers/Cargo.toml b/crates/vendor/candle/candle-transformers/Cargo.toml index e7194aa624..e9786e471a 100644 --- a/crates/vendor/candle/candle-transformers/Cargo.toml +++ b/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 = [] diff --git a/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/lcm.rs b/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/lcm.rs index 96fbc2cdfd..6f2a780859 100644 --- a/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/lcm.rs +++ b/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/lcm.rs @@ -1,6 +1,6 @@ use super::schedulers::{betas_for_alpha_bar, BetaSchedule, PredictionType}; use candle::{DType, Device, IndexOp, Result, Tensor}; - +const FINAL_ALPHA_CUMPROD: f64 = 0.9991499781608582; #[derive(Debug, Clone, PartialEq, Eq)] pub enum LCMVarianceType { FixedSmall, @@ -60,7 +60,7 @@ pub struct LCMScheduler { } impl LCMScheduler { - pub fn new(inference_steps: usize, config: LCMSchedulerConfig) -> Result { + pub fn new(inference_steps: usize, strength: f64, config: LCMSchedulerConfig) -> Result { let betas = match config.beta_schedule { BetaSchedule::ScaledLinear => super::utils::linspace( config.beta_start.sqrt(), @@ -84,7 +84,6 @@ impl LCMScheduler { } // Create proper LCM timesteps to match diffusers - let strength = 0.8; let timesteps = LCMScheduler::get_timesteps_for_steps(inference_steps, strength); let step_ratio = config.train_timesteps / inference_steps; @@ -335,7 +334,7 @@ impl LCMScheduler { let alpha_prod_t = self.alphas_cumprod[timestep]; let alpha_prod_t_prev = if prev_timestep == 0 { - 1.0 + FINAL_ALPHA_CUMPROD } else { self.alphas_cumprod[prev_timestep] }; @@ -355,7 +354,7 @@ impl LCMScheduler { } }; - let denoised = ((pred_original_sample * c_out)? + (sample * c_skip)?)?; + let denoised = ((c_out * pred_original_sample )? + (c_skip * sample)?)?; let pred_prev_sample = if step_index != self.timesteps.len() - 1 { println!("Adding noise at intermediate step {}", timestep); @@ -470,6 +469,7 @@ impl LCMScheduler { // } pub fn add_noise(&self, original: &Tensor, noise: Tensor, timestep: usize) -> Result { + println!("Adding noise at timestep {}", timestep); let timestep = if timestep >= self.alphas_cumprod.len() { self.alphas_cumprod.len() - 1 // Ensure valid range } else {