diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index 6c6f6634dd..8df4795ceb 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs @@ -14,12 +14,15 @@ use tauri::State; const PROMPT_FILENAME : &str = "prompt.txt"; const PROMPT: &str = "A beautiful landscape with mountains and a lake"; +const NEGATIVE_PROMPT: &str = "bad quality, bad faces, poor quality, blurry faces, watermark"; + /// This handler takes an image (as a base64 encoded string) and a prompt and returns /// an image (as a base64-encoded string). #[tauri::command] pub fn infer_image( image: &str, prompt: Option, + strength: Option, model_config: State, model_cache: State, prompt_cache: State, @@ -32,7 +35,7 @@ pub fn infer_image( let image = hydrate_base64_image(image) .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?; - let result = do_infer_image(&prompt, image, &model_config, &model_cache, prompt_cache); + let result = do_infer_image(&prompt, image, strength, &model_config, &model_cache, prompt_cache); if let Err(err) = result.as_deref() { error!("There was an error: {:?}", err); @@ -44,6 +47,7 @@ pub fn infer_image( fn do_infer_image( prompt: &str, image: DynamicImage, + strength: Option, config: &AppConfig, model_cache: &ModelCache, prompt_cache: State, @@ -53,12 +57,13 @@ fn do_infer_image( let args = Args { image: &image, prompt: prompt.to_string(), - uncond_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, }; match stable_diffusion_pipeline(args) { 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 9ac4a390ef..0072d021a6 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 @@ -23,6 +23,7 @@ pub struct Args<'a> { 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 { @@ -96,7 +97,11 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { println!("Initial latents shape: {:?}", latents.shape()); println!("Calculating start step for diffusion process..."); - let img2img_strength = 0.75f64; + + let img2img_strength = match args.strength { + None => 0.75f64, + Some(strength) => (strength as f64) / 100.0f64, + }; let t_start = { let start = args.configs.scheduler_steps - (args.configs.scheduler_steps as f64 * img2img_strength) as usize;