diff --git a/Cargo.lock b/Cargo.lock index 7b83a0e7a2..666ff764fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -658,6 +658,7 @@ name = "app" version = "0.1.0" dependencies = [ "anyhow", + "assertor", "base64 0.22.1", "bytes 1.10.0", "candle-core", @@ -676,6 +677,7 @@ dependencies = [ "log", "memory_caching", "once_cell", + "primitives", "rand 0.8.5", "reqwest 0.12.12", "serde", @@ -856,6 +858,15 @@ version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "10d2119f741b79fe9907f5396d19bffcb46568cfcc315e78677d731972ac7085" +[[package]] +name = "assertor" +version = "0.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b02f2af4042fbbed4d5732c517666629512c574c8e2edb5cfe2cf8a86564d52" +dependencies = [ + "num-traits", +] + [[package]] name = "async-channel" version = "1.9.0" @@ -8381,6 +8392,7 @@ dependencies = [ name = "primitives" version = "0.0.1" dependencies = [ + "assertor", "my-workspace-hack", "testing", ] diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index 919591e7e3..9275426daf 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -22,6 +22,7 @@ tauri-build = { version = "2.0.4", features = [] } # Internal filesys = { path = "../../lib/files/filesys" } memory_caching = { path = "../../lib/caching/memory_caching" } +primitives = { path = "../../lib/data/primitives" } # Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332 cudarc = "0.13.7" @@ -71,6 +72,7 @@ cuda = [ [dev-dependencies] libloading = "0.8.6" serial_test = "3.2.0" +assertor = "0.0.3" [profile.dev] incremental = true # Compile your binary in smaller steps. diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index 3eb1fdc90b..506874bd42 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs @@ -26,7 +26,6 @@ const RANDOM_SEED: u32 = 42; pub async fn infer_image( image: &str, prompt: Option, - // NB: `strength` is supposed to be a u8 in range 0-100, but Javascript keeps sending us floats (eg. 14.000008) strength: Option, model_config: State<'_, AppConfig>, model_cache: State<'_, ModelCache>, @@ -42,8 +41,6 @@ pub async fn infer_image( let image = hydrate_base64_image(image) .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?; - let strength = strength.map(|s| float_to_u8(s)); - let result = do_infer_image( &prompt, image, @@ -65,7 +62,7 @@ pub async fn infer_image( async fn do_infer_image( prompt: &str, image: DynamicImage, - strength: Option, + strength: Option, config: &AppConfig, model_cache: &ModelCache, prompt_cache: State<'_, PromptCache>, 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 a7deeed19b..dfed61b3f6 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 @@ -9,6 +9,7 @@ use crate::ml::model_file::{ModelFile, 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::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range; use crate::ml::weights_registry::weights::{LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; use crate::state::app_config::AppConfig; use crate::state::app_dir::AppDataRoot; @@ -25,12 +26,14 @@ use log::info; use rand::Rng; use tauri::{AppHandle, Emitter}; +const DEFAULT_STRENGTH : f64 = 75.0; + pub struct Args<'a> { pub image: &'a DynamicImage, pub prompt: String, pub uncond_prompt: String, pub cfg_scale: Option, - pub i2i_strength: Option, + pub i2i_strength: Option, pub configs: &'a AppConfig, pub model_cache: &'a ModelCache, pub prompt_cache: &'a PromptCache, @@ -43,8 +46,13 @@ 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, use_flash_attn } = args; let weights_dir = app_data_root.weights_dir(); + + let img2img_strength = args.i2i_strength + .map(|s| remap_lcm_strength_range(s)) + .unwrap_or(DEFAULT_STRENGTH); + + let img2img_strength = img2img_strength / 100.0; - 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 @@ -181,12 +189,6 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result { 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)?; diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/mod.rs b/crates/desktop/src-tauri/src/ml/stable_diffusion/mod.rs index 67e6759bd0..d605b647d8 100644 --- a/crates/desktop/src-tauri/src/ml/stable_diffusion/mod.rs +++ b/crates/desktop/src-tauri/src/ml/stable_diffusion/mod.rs @@ -1,4 +1,5 @@ pub mod get_vae_scale; pub mod infer_clip_text_embeddings; pub mod stable_diffusion_pipeline; -pub mod lcm_pipeline; \ No newline at end of file +pub mod lcm_pipeline; +mod remap_lcm_strength_range; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/remap_lcm_strength_range.rs b/crates/desktop/src-tauri/src/ml/stable_diffusion/remap_lcm_strength_range.rs new file mode 100644 index 0000000000..dd9f248251 --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/stable_diffusion/remap_lcm_strength_range.rs @@ -0,0 +1,54 @@ +use primitives::numerics::remap_range_f64::{remap_range_f64, Args}; + +/// Our LCM implementation uses strength [13, 100]. Strength below 13 will take 8 steps, +/// so we need to limit ourselves to this range from an input range of [0, 100]. +pub fn remap_lcm_strength_range(value: f64) -> f64 { + remap_range_f64(Args { + old_min: 0.0, + old_max: 100.0, + new_min: 13.0, + new_max: 100.0, + value, + }) +} + +#[cfg(test)] +mod test { + use crate::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range; + use assertor::{assert_that, FloatAssertion}; + + #[test] + fn old_max() { + assert_eq!(remap_lcm_strength_range(100.0), 100.0); + } + + #[test] + fn old_min() { + assert_eq!(remap_lcm_strength_range(0.0), 13.0); + } + + #[test] + fn under_bounds() { + assert_eq!(remap_lcm_strength_range(-100.0), 13.0); + } + + #[test] + fn over_bounds() { + assert_eq!(remap_lcm_strength_range(200.0), 100.0); + } + + #[test] + fn other_values() { + assert_that!(remap_lcm_strength_range(0.0)).with_abs_tol(0.001).is_approx_equal_to(13.000); + assert_that!(remap_lcm_strength_range(5.0)).with_abs_tol(0.001).is_approx_equal_to(17.350); + assert_that!(remap_lcm_strength_range(11.0)).with_abs_tol(0.001).is_approx_equal_to(22.570); + assert_that!(remap_lcm_strength_range(13.0)).with_abs_tol(0.001).is_approx_equal_to(24.310); + assert_that!(remap_lcm_strength_range(14.0)).with_abs_tol(0.001).is_approx_equal_to(25.180); + assert_that!(remap_lcm_strength_range(20.0)).with_abs_tol(0.001).is_approx_equal_to(30.400); + assert_that!(remap_lcm_strength_range(25.0)).with_abs_tol(0.001).is_approx_equal_to(34.750); + assert_that!(remap_lcm_strength_range(50.0)).with_abs_tol(0.001).is_approx_equal_to(56.500); + assert_that!(remap_lcm_strength_range(75.0)).with_abs_tol(0.001).is_approx_equal_to(78.250); + assert_that!(remap_lcm_strength_range(99.0)).with_abs_tol(0.001).is_approx_equal_to(99.130); + assert_that!(remap_lcm_strength_range(100.0)).with_abs_tol(0.001).is_approx_equal_to(100.000); + } +} diff --git a/crates/lib/data/primitives/Cargo.toml b/crates/lib/data/primitives/Cargo.toml index 97a0e27a00..a9ed00d3be 100644 --- a/crates/lib/data/primitives/Cargo.toml +++ b/crates/lib/data/primitives/Cargo.toml @@ -32,3 +32,4 @@ my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" } [dev-dependencies] testing = { path = "../../testing" } +assertor = "0.0.3" diff --git a/crates/lib/data/primitives/src/numerics/mod.rs b/crates/lib/data/primitives/src/numerics/mod.rs index 0a327df77d..d4b2921c49 100644 --- a/crates/lib/data/primitives/src/numerics/mod.rs +++ b/crates/lib/data/primitives/src/numerics/mod.rs @@ -1,5 +1,6 @@ pub mod i32_to_u32_zero_clamped; pub mod i64_to_u64_zero_clamped; +pub mod remap_range_f64; pub mod u64_to_i32_saturating; pub mod u64_to_i64_saturating; pub mod u64_to_u32_saturating; \ No newline at end of file diff --git a/crates/lib/data/primitives/src/numerics/remap_range_f64.rs b/crates/lib/data/primitives/src/numerics/remap_range_f64.rs new file mode 100644 index 0000000000..19a48adf67 --- /dev/null +++ b/crates/lib/data/primitives/src/numerics/remap_range_f64.rs @@ -0,0 +1,131 @@ + +pub struct Args { + pub old_min: f64, + pub old_max: f64, + pub new_min: f64, + pub new_max: f64, + + /// The value to remap to the new range. + pub value: f64, +} + +pub fn remap_range_f64(args: Args) -> f64 { + let old_range = args.old_max - args.old_min; + let new_range = args.new_max - args.new_min; + + let new_value = (args.value - args.old_min) * new_range / old_range + args.new_min; + + if new_value > args.new_max { + args.new_max + } else if new_value < args.new_min { + args.new_min + } else { + new_value + } +} + +#[cfg(test)] +mod tests { + use crate::numerics::remap_range_f64::{remap_range_f64, Args}; + use assertor::{assert_that, FloatAssertion}; + + #[test] + fn test_base_cases() { + assert_that!(remap_range_f64(Args { + old_min: 0.0, + old_max: 1.0, + new_min: 0.0, + new_max: 1.0, + value: 0.0, + })).with_abs_tol(0.01).is_approx_equal_to(0.0); + + assert_that!(remap_range_f64(Args { + old_min: 0.0, + old_max: 1.0, + new_min: 0.0, + new_max: 1.0, + value: 1.0, + })).with_abs_tol(0.01).is_approx_equal_to(1.0); + } + + #[test] + fn test_maxima() { + assert_that!(remap_range_f64(Args { + old_min: 0.0, + old_max: 100.0, + new_min: 0.0, + new_max: 1.0, + value: 100.0, + })).with_abs_tol(0.01).is_approx_equal_to(1.0); + + assert_that!(remap_range_f64(Args { + old_min: 0.0, + old_max: 1.0, + new_min: 0.0, + new_max: 50.0, + value: 1.0, + })).with_abs_tol(0.01).is_approx_equal_to(50.0); + + assert_that!(remap_range_f64(Args { + old_min: 0.0, + old_max: 1.0, + new_min: -100.0, + new_max: 100.0, + value: 1.0, + })).with_abs_tol(0.01).is_approx_equal_to(100.0); + } + + #[test] + fn test_minima() { + assert_that!(remap_range_f64(Args { + old_min: 0.0, + old_max: 100.0, + new_min: 0.0, + new_max: 1.0, + value: 0.0, + })).with_abs_tol(0.01).is_approx_equal_to(0.0); + + assert_that!(remap_range_f64(Args { + old_min: 0.0, + old_max: 1.0, + new_min: 0.0, + new_max: 50.0, + value: 0.0, + })).with_abs_tol(0.01).is_approx_equal_to(0.0); + + assert_that!(remap_range_f64(Args { + old_min: 0.0, + old_max: 1.0, + new_min: -100.0, + new_max: 100.0, + value: 0.0, + })).with_abs_tol(0.01).is_approx_equal_to(-100.0); + } + + #[test] + fn test_other_values() { + assert_that!(remap_range_f64(Args { + old_min: 0.0, + old_max: 100.0, + new_min: 0.0, + new_max: 1.0, + value: 50.0, + })).with_abs_tol(0.01).is_approx_equal_to(0.50); + + assert_that!(remap_range_f64(Args { + old_min: 0.0, + old_max: 1.0, + new_min: 0.0, + new_max: 50.0, + value: 0.10, + })).with_abs_tol(0.01).is_approx_equal_to(5.0); + + assert_that!(remap_range_f64(Args { + old_min: 0.0, + old_max: 1.0, + new_min: -100.0, + new_max: 100.0, + value: 0.25, + })).with_abs_tol(0.01).is_approx_equal_to(-50.0); + } +}