diff --git a/crates/desktop/src-tauri/binary_includes/1024.png b/crates/desktop/src-tauri/binary_includes/1024.png new file mode 100644 index 0000000000..9f746687dc Binary files /dev/null and b/crates/desktop/src-tauri/binary_includes/1024.png differ diff --git a/crates/desktop/src-tauri/src/endpoints/mod.rs b/crates/desktop/src-tauri/src/endpoints/mod.rs index d365df449b..f49bdce06e 100644 --- a/crates/desktop/src-tauri/src/endpoints/mod.rs +++ b/crates/desktop/src-tauri/src/endpoints/mod.rs @@ -1,5 +1,6 @@ pub mod download_models; pub mod flip_image; -pub mod image_endpoint; +pub mod realtime_image_endpoint; pub mod remove_background_endpoint; pub mod test_counter; +pub mod text_to_image_endpoint; diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs similarity index 100% rename from crates/desktop/src-tauri/src/endpoints/image_endpoint.rs rename to crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs diff --git a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs new file mode 100644 index 0000000000..081625fb81 --- /dev/null +++ b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs @@ -0,0 +1,99 @@ +use crate::ml::model_cache::ModelCache; +use crate::ml::prompt_cache::PromptCache; +use crate::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}; +use crate::state::app_config::AppConfig; +use crate::state::app_dir::AppDataRoot; +use crate::utils::image::decode_base64_image::decode_base64_image; +use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png; +use image::imageops::FilterType; +use image::{DynamicImage, ImageReader, RgbImage}; +use log::{error, info}; +use once_cell::sync::Lazy; +use std::io::Cursor; +use std::path::PathBuf; +use tauri::{AppHandle, State}; + +const RANDOM_SEED: u32 = 42; + +const STRENGTH: f64 = 1.0; +const PNG_BYTES : &[u8] = include_bytes!("../../binary_includes/1024.png"); + +// NB: We're just hacking a temporary "text to image" test endpoint leveraging the existing image to image modality. +static IMAGE : Lazy = Lazy::new(||{ + ImageReader::new(Cursor::new(PNG_BYTES)) + .with_guessed_format() + .expect("failed to detect png format") + .decode() + .expect("failed to decode image") +}); + +#[tauri::command] +pub async fn text_to_image( + prompt: String, + model_config: State<'_, AppConfig>, + model_cache: State<'_, ModelCache>, + prompt_cache: State<'_, PromptCache>, + app_data_root: State<'_, AppDataRoot>, + app: AppHandle, +) -> Result { + info!("text_to_image endpoint called."); + + let image = IMAGE.resize(512, 512, FilterType::CatmullRom); + + let result = do_infer_image( + &prompt, + image, + Some(STRENGTH), + &model_config, + &model_cache, + prompt_cache, + app, + &app_data_root + ).await; + + let image = match result { + Ok(image) => image, + Err(err) => { + error!("There was an error generating the image: {:?}", err); + return Err(format!("There was an error generating the image: {}", err)); + } + }; + + let bytes = encode_rgb_image_base64_png(image) + .map_err(|err| format!("failure to encode image: {:?}", err))?; + + info!("Inference successful; image converted to base64, serving back to browser..."); + + Ok(bytes) +} + +async fn do_infer_image( + prompt: &str, + image: DynamicImage, + strength: Option, + config: &AppConfig, + model_cache: &ModelCache, + prompt_cache: State<'_, PromptCache>, + app: AppHandle, + app_data_root: &AppDataRoot, +) -> Result { + + let args = Args { + image: &image, + prompt: prompt.to_string(), + uncond_prompt: "".to_string(), + model_cache, + configs: config, + prompt_cache: &prompt_cache, + i2i_strength: strength, + cfg_scale: config.cfg_scale, + app: &app, + app_data_root, + use_flash_attn: true, + }; + + let image = lcm_pipeline(args) + .map_err(|err| format!("failure to encode image: {:?}", err))?; + + Ok(image) +} diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index a5f00e13e2..653655c2a0 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -7,9 +7,10 @@ pub mod transfer; pub mod utils; use crate::endpoints::download_models::download_models; -use crate::endpoints::image_endpoint::infer_image; +use crate::endpoints::realtime_image_endpoint::infer_image; use crate::endpoints::remove_background_endpoint::remove_background; use crate::endpoints::test_counter::test_counter; +use crate::endpoints::text_to_image_endpoint::text_to_image; use crate::ml::model_cache::ModelCache; use crate::ml::prompt_cache::PromptCache; use crate::state::app_config::AppConfig; @@ -66,6 +67,7 @@ pub fn run() { infer_image, remove_background, test_counter, + text_to_image, ]) .run(tauri::generate_context!()) .expect("error while running tauri application"); diff --git a/crates/desktop/src-tauri/tauri.conf.json b/crates/desktop/src-tauri/tauri.conf.json index 718a24832a..17c4b70706 100644 --- a/crates/desktop/src-tauri/tauri.conf.json +++ b/crates/desktop/src-tauri/tauri.conf.json @@ -32,9 +32,6 @@ "icons/128x128@2x.png", "icons/icon.icns", "icons/icon.ico" - ], - "resources": { - "../../../windows_dlls/*.dll": "." - } + ] } }