diff --git a/app_config.yaml b/app_config.yaml index 924b3516dc..fca2ac4c14 100644 --- a/app_config.yaml +++ b/app_config.yaml @@ -4,4 +4,4 @@ image_height: 512 # NB(bt): It looks like we can get down to 2 sampling steps, perhaps ? scheduler_steps: 2 scheduler_samples: 1 - +cfg_scale: 1.1 diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index 8df4795ceb..3b958e9e55 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs @@ -59,11 +59,11 @@ fn do_infer_image( prompt: prompt.to_string(), uncond_prompt: NEGATIVE_PROMPT.to_string(), //guidance_scale: Some(0.0), - guidance_scale: None, model_cache, configs: config, prompt_cache: &prompt_cache, - strength: strength, + i2i_strength: strength, + cfg_scale: config.cfg_scale, }; match stable_diffusion_pipeline(args) { 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 eed7961c00..8a4836f729 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 @@ -48,7 +48,7 @@ pub fn infer_clip_text_embeddings( false, // use_f16 &device, dtype, - false, // use_guide_scale + use_guide_scale, *first, ) }) diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs b/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs index 0072d021a6..9cbdad1ef7 100644 --- a/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs +++ b/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs @@ -3,27 +3,27 @@ use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}; 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::prompt_cache::PromptCache; 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::{Error as E, Result}; -use candle_core::{DType, IndexOp, D}; +use anyhow::{anyhow, Error as E, Result}; +use candle_core::{DType, IndexOp, Tensor, D}; use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution; use image::DynamicImage; use log::info; use rand::Rng; - pub struct Args<'a> { pub image: &'a DynamicImage, pub prompt: String, pub uncond_prompt: String, - pub guidance_scale: Option, + pub cfg_scale: Option, + pub i2i_strength: Option, pub configs: &'a AppConfig, pub model_cache: &'a ModelCache, pub prompt_cache: &'a PromptCache, - pub strength: Option, } pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { @@ -38,20 +38,33 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { // TODO(bt,2025-02-18): The scheduler is `EulerAncestralDiscreteScheduler`, but we may want to port an LCM scheduler. // This is a target for performance improvement // See: https://github.com/huggingface/candle/issues/1331 - println!("Building scheduler..."); let mut scheduler = args.configs.sd_config.build_scheduler( args.configs.scheduler_steps)?; let seed = args.configs.seed.unwrap_or_else(|| rand::thread_rng().gen()); + args.configs.device.set_seed(seed)?; println!("Using seed: {}", seed); - args.configs.device.set_seed(seed)?; + + let guidance_scale = match args.cfg_scale { + Some(guidance_scale) => guidance_scale, + None => match args.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; + + println!("Using guide scale: {}", use_guide_scale); info!("Checking if prompt is cached"); let maybe_cached = args.prompt_cache.get_copy(&args.prompt)?; - - let text_embeddings = if let Some(tensor) = maybe_cached { - info!("Prompt is cached!"); + + let mut text_embeddings = if let Some(tensor) = maybe_cached { tensor } else { info!("Prompt is NOT cached! Calculating embedding..."); @@ -66,12 +79,12 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { false, // use_f16 &args.configs.device, args.configs.dtype, - false, // use_guide_scale + use_guide_scale, )?; args.prompt_cache.store_copy(&args.prompt, &tensor)?; tensor }; - + println!("Text embeddings shape: {:?}", text_embeddings.shape()); println!("Loading input image into tensor..."); @@ -97,8 +110,8 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { println!("Initial latents shape: {:?}", latents.shape()); println!("Calculating start step for diffusion process..."); - - let img2img_strength = match args.strength { + + let img2img_strength = match args.i2i_strength { None => 0.75f64, Some(strength) => (strength as f64) / 100.0f64, }; @@ -136,20 +149,29 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { } println!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep); - - // TODO(bt,2025-08-18): Do we have to clone the tensor? - let latent_model_input = scheduler.scale_model_input(latents.clone(), timestep)?; - - let noise_pred = match args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings) { - Ok(pred) => { - pred - }, + + let latent_model_input = if use_guide_scale { + Tensor::cat(&[&latents, &latents], 0)? + } else { + latents.clone() // TODO(bt,2025-08-18): Do we have to clone the tensor? `scale_model_input` requires ownership + }; + + let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?; + + let mut noise_pred = match args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings) { + 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)?)?)?; + } + latents = scheduler.step(&noise_pred, timestep, &latents)?; } diff --git a/crates/desktop/src-tauri/src/state/app_config.rs b/crates/desktop/src-tauri/src/state/app_config.rs index 3c5f5e0a7c..0d22f1c5fd 100644 --- a/crates/desktop/src-tauri/src/state/app_config.rs +++ b/crates/desktop/src-tauri/src/state/app_config.rs @@ -32,6 +32,8 @@ pub struct AppConfig { pub scheduler_samples: usize, pub seed: Option, + + pub cfg_scale: Option, /// Probably shouldn't be here. pub hf_api: Api, @@ -72,6 +74,7 @@ impl AppConfig { scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(1), scheduler_samples: yaml_configs.scheduler_samples.unwrap_or(15), seed: yaml_configs.seed, + cfg_scale: yaml_configs.cfg_scale, hf_api, }) } diff --git a/crates/desktop/src-tauri/src/state/yaml_config.rs b/crates/desktop/src-tauri/src/state/yaml_config.rs index cd0626338c..35037c5eb5 100644 --- a/crates/desktop/src-tauri/src/state/yaml_config.rs +++ b/crates/desktop/src-tauri/src/state/yaml_config.rs @@ -12,6 +12,7 @@ pub struct YamlConfig { pub scheduler_steps: Option, pub scheduler_samples: Option, pub seed: Option, + pub cfg_scale: Option, } impl YamlConfig {