diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index d8522b802d..a6167c44bb 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs @@ -1,13 +1,13 @@ use crate::endpoints::sd::{run, Args}; -use crate::state::yaml_config::YamlConfig; use crate::ml::model_cache::ModelCache; use crate::ml::prompt_cache::PromptCache; use crate::state::app_config::AppConfig; +use crate::state::yaml_config::YamlConfig; use anyhow::Result; use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; use image::imageops::FilterType; -use image::ImageReader; -use log::error; +use image::{DynamicImage, EncodableLayout, ImageReader}; +use log::{error, info}; use std::io::Cursor; use std::path::PathBuf; use tauri::State; @@ -46,13 +46,7 @@ pub fn infer_image( // TODO(bt,2025-02-17): Running out of vram with full image buffer size let image = image.resize(512, 512, FilterType::CatmullRom); - // TODO(bt,2025-02-14): Use byte buffers instead of hitting the filesystem. - let image_path = PathBuf::from("input_image.png"); - - image.save(&image_path) - .map_err(|err| format!("Failed to save image: {}", err))?; - - let result = do_infer_image(&prompt, &image_path, &model_config, &model_cache, prompt_cache); + let result = do_infer_image(&prompt, image, &model_config, &model_cache, prompt_cache); if let Err(err) = result.as_deref() { error!("There was an error: {:?}", err); @@ -63,7 +57,7 @@ pub fn infer_image( fn do_infer_image( prompt: &str, - image_path: &PathBuf, + image: DynamicImage, config: &AppConfig, model_cache: &ModelCache, prompt_cache: State, @@ -71,7 +65,7 @@ fn do_infer_image( println!("infer_image called; generating image with SDXL Turbo..."); let args = Args { - image_path, + image: &image, prompt: prompt.to_string(), uncond_prompt: "".to_string(), //guidance_scale: Some(0.0), @@ -82,17 +76,27 @@ fn do_infer_image( }; match run(args) { - Ok(_) => { + Ok(image) => { + //info!("Image len: {:?}", image.len()); + //let bytes = BASE64_STANDARD.encode(image.as_bytes()); + //// [2025-02-18][04:14:15][app_lib::endpoints::image_endpoint][INFO] First bytes: "DAwODAwNDA" + //info!("First bytes: {:?}", bytes.split_at(10).0); + + let img_data = std::fs::read("temp.png") .map_err(|e| format!("Failed to read generated image: {}", e))?; - - let encoded = BASE64_STANDARD.encode(&img_data); + + let bytes = BASE64_STANDARD.encode(&img_data); println!("Generated image encoded successfully"); let _ = std::fs::remove_file("temp.png"); - Ok(encoded) - } - Err(e) => Err(format!("Failed to generate image: {}", e)) + + + + + Ok(bytes) + }, + Err(e) => Err(format!("Failed to generate image: {}", e)), } } diff --git a/crates/desktop/src-tauri/src/endpoints/sd.rs b/crates/desktop/src-tauri/src/endpoints/sd.rs index e5c6d0c18c..86b76bbece 100644 --- a/crates/desktop/src-tauri/src/endpoints/sd.rs +++ b/crates/desktop/src-tauri/src/endpoints/sd.rs @@ -1,37 +1,27 @@ -#[cfg(feature = "accelerate")] -extern crate accelerate_src; -#[cfg(feature = "mkl")] -extern crate intel_mkl_src; - -use std::mem::forget; -use std::path::Path; -use std::sync::{Arc, LockResult, RwLock, RwLockReadGuard, RwLockWriteGuard}; -use candle_transformers::models::stable_diffusion; +use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}; +use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor; +use crate::ml::image::save_image_from_tensor::save_image_from_tensor; +use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage}; use crate::ml::infer_clip_text_embeddings::infer_clip_text_embeddings; -use crate::ml::model_file::StableDiffusionVersion; -use crate::ml::save_image_from_tensor::save_image_from_tensor; -use anyhow::{Error as E, Result}; -use candle_core::{DType, Device, IndexOp, Module, Tensor, D}; -use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution; -use hf_hub::api::sync::Api; -use log::info; -use once_cell::sync::Lazy; -use rand::Rng; -use crate::ml::create_inpainting_tensors::create_inpainting_tensors; -use crate::ml::load_image_file_to_tensor::load_image_file_to_tensor; -use crate::ml::load_image_file_to_tensor_2::load_image_file_to_tensor_2; use crate::ml::model_cache::ModelCache; -use crate::ml::models::unet_model::UNetModel; use crate::ml::prompt_cache::PromptCache; +use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale; use crate::state::app_config::AppConfig; +use anyhow::{Error as E, Result}; +use candle_core::{DType, Device, IndexOp, D}; +use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution; +use image::DynamicImage; +use log::info; +use rand::Rng; + // TODO: Clean up // Note for Kasisnu: I'm going to start using lifetimes as long as that doesn't slow your velocity // Basically the args o this telescopic args struct are guaranteed to live as long as the struct // itself with the 'a lifetime. pub struct Args<'a> { - pub image_path: &'a Path, + pub image: &'a DynamicImage, pub prompt: String, pub uncond_prompt: String, pub guidance_scale: Option, @@ -40,7 +30,7 @@ pub struct Args<'a> { pub prompt_cache: &'a PromptCache, } -pub fn run(args: Args<'_>) -> Result<()> { +pub fn run(args: Args<'_>) -> Result { println!("Starting image generation with the following configuration:"); println!(" Model: {:?}", args.configs.sd_version); println!(" Prompt: {}", args.prompt); @@ -119,17 +109,14 @@ pub fn run(args: Args<'_>) -> Result<()> { println!("Loading input image into tensor..."); - let input_image = load_image_file_to_tensor_2(args.image_path)?; - - let input_image = input_image.to_device(&args.configs.device)? - .to_dtype(args.configs.dtype)?; + let input_image = dynamic_image_to_tensor( + args.image, + &args.configs.device, + args.configs.dtype)?; // REFERENCE IMAGE shape: [1, 3, 1024, 1024] println!("REFERENCE IMAGE shape: {:?}", input_image.shape()); - - - println!("Calculating start step for diffusion process..."); let img2img_strength = 0.75f64; @@ -142,17 +129,7 @@ pub fn run(args: Args<'_>) -> Result<()> { let use_guide_scale = false; // TODO - let vae_scale = match args.configs.sd_version { - StableDiffusionVersion::V1_5 - | StableDiffusionVersion::V1_5Inpaint - | StableDiffusionVersion::V2_1 - | StableDiffusionVersion::V2Inpaint - | StableDiffusionVersion::XlInpaint - | StableDiffusionVersion::Xl => 0.18215, - StableDiffusionVersion::Turbo => 0.13025, - }; - - + let vae_scale = get_vae_scale(args.configs.sd_version); let encoded_image : DiagonalGaussianDistribution = args.model_cache.vae_encode(&input_image)?; let init_latent_dist = encoded_image; @@ -173,15 +150,6 @@ pub fn run(args: Args<'_>) -> Result<()> { println!("Latents initialized successfully"); - - - let vae_scale = match args.configs.sd_version { - StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint | StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Xl => 0.18215, - StableDiffusionVersion::Turbo => 0.13025, - }; - - - //let encoded_sample = encoded_image.sample()?; // Encoded sample shape: [1, 4, 128, 128] ... so close. Just needs to be smaller. @@ -190,8 +158,6 @@ pub fn run(args: Args<'_>) -> Result<()> { //let mask_latents = (encoded_sample * vae_scale)?.to_device(&args.model_configs.device)?; //println!("Mask latents shape: {:?}", mask_latents.shape()); - - //let timesteps = scheduler.timesteps().to_vec(); let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect(); @@ -321,8 +287,12 @@ pub fn run(args: Args<'_>) -> Result<()> { let image = (image.clamp(0f32, 1.)? * 255.)?.to_dtype(DType::U8)?; println!("Converted to 8-bit format"); - save_image_from_tensor(&image.i(0)?, "temp.png")?; - println!("Image generation completed successfully"); + //save_image_from_tensor(&image.i(0)?, "temp.png")?; + //println!("Image generation completed successfully"); + + let image = tensor_to_image_buffer(&image.i(0)?)?; - Ok(()) + image.save("temp.png").map_err(E::from)?; + + Ok(image) } diff --git a/crates/desktop/src-tauri/src/ml/image/dynamic_image_to_tensor.rs b/crates/desktop/src-tauri/src/ml/image/dynamic_image_to_tensor.rs new file mode 100644 index 0000000000..47d83deda9 --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/image/dynamic_image_to_tensor.rs @@ -0,0 +1,17 @@ +use anyhow::Result; +use candle_core::{DType, Device, Tensor}; +use image::DynamicImage; + +pub fn dynamic_image_to_tensor(image: &DynamicImage, device: &Device, datatype: DType) -> Result { + let (height, width) = (image.height() as usize, image.width() as usize); + let height = height - height % 32; + let width = width - width % 32; + let image = image.resize_to_fill(width as u32, height as u32, image::imageops::FilterType::CatmullRom); + let image = image.to_rgb8(); + let image = image.into_raw(); + Ok(Tensor::from_vec(image, (height, width, 3), device)? + .permute((2, 0, 1))? + .to_dtype(datatype)? + .affine(2. / 255., -1.)? + .unsqueeze(0)?) +} diff --git a/crates/desktop/src-tauri/src/ml/load_image_file_to_tensor.rs b/crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/load_image_file_to_tensor.rs rename to crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor.rs diff --git a/crates/desktop/src-tauri/src/ml/load_image_file_to_tensor_2.rs b/crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor_2.rs similarity index 80% rename from crates/desktop/src-tauri/src/ml/load_image_file_to_tensor_2.rs rename to crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor_2.rs index 3b50aed4c5..6eb5ff4aea 100644 --- a/crates/desktop/src-tauri/src/ml/load_image_file_to_tensor_2.rs +++ b/crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor_2.rs @@ -1,10 +1,3 @@ -#[cfg(feature = "accelerate")] -extern crate accelerate_src; -#[cfg(feature = "mkl")] -extern crate intel_mkl_src; - -use candle_transformers::models::stable_diffusion; - use anyhow::Result; use candle_core::{DType, Device, Tensor}; diff --git a/crates/desktop/src-tauri/src/ml/image/mod.rs b/crates/desktop/src-tauri/src/ml/image/mod.rs new file mode 100644 index 0000000000..b6059fd6e3 --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/image/mod.rs @@ -0,0 +1,5 @@ +pub mod dynamic_image_to_tensor; +pub mod load_image_file_to_tensor; +pub mod load_image_file_to_tensor_2; +pub mod save_image_from_tensor; +pub mod tensor_to_image_buffer; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/ml/save_image_from_tensor.rs b/crates/desktop/src-tauri/src/ml/image/save_image_from_tensor.rs similarity index 89% rename from crates/desktop/src-tauri/src/ml/save_image_from_tensor.rs rename to crates/desktop/src-tauri/src/ml/image/save_image_from_tensor.rs index a9917b385e..145ac69e50 100644 --- a/crates/desktop/src-tauri/src/ml/save_image_from_tensor.rs +++ b/crates/desktop/src-tauri/src/ml/image/save_image_from_tensor.rs @@ -1,8 +1,3 @@ -#[cfg(feature = "accelerate")] -extern crate accelerate_src; -#[cfg(feature = "mkl")] -extern crate intel_mkl_src; - use anyhow::{Error as E, Result}; use candle_core::Tensor; use image::{ImageBuffer, Rgb}; diff --git a/crates/desktop/src-tauri/src/ml/image/tensor_to_image_buffer.rs b/crates/desktop/src-tauri/src/ml/image/tensor_to_image_buffer.rs new file mode 100644 index 0000000000..362f954b24 --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/image/tensor_to_image_buffer.rs @@ -0,0 +1,18 @@ +use candle_core::Tensor; +use image::{ImageBuffer, Rgb}; + +pub type RgbImage = ImageBuffer, Vec>; + +pub fn tensor_to_image_buffer(img: &Tensor) -> anyhow::Result { + let (channel, height, width) = img.dims3()?; + + if channel != 3 { + anyhow::bail!("save_image expects an input of shape (3, height, width)") + } + + let img = img.permute((1, 2, 0))?.flatten_all()?; + let pixels = img.to_vec1::()?; + + ImageBuffer::from_raw(width as u32, height as u32, pixels) + .ok_or_else(|| anyhow::anyhow!("error loading image from raw tensor")) +} diff --git a/crates/desktop/src-tauri/src/ml/mod.rs b/crates/desktop/src-tauri/src/ml/mod.rs index 9fd2560f35..3f95f0499e 100644 --- a/crates/desktop/src-tauri/src/ml/mod.rs +++ b/crates/desktop/src-tauri/src/ml/mod.rs @@ -1,9 +1,8 @@ pub mod create_inpainting_tensors; +pub mod image; pub mod infer_clip_text_embeddings; -pub mod load_image_file_to_tensor; -pub mod load_image_file_to_tensor_2; pub mod model_cache; pub mod model_file; pub mod models; pub mod prompt_cache; -pub mod save_image_from_tensor; \ No newline at end of file +pub mod stable_diffusion; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/get_vae_scale.rs b/crates/desktop/src-tauri/src/ml/stable_diffusion/get_vae_scale.rs new file mode 100644 index 0000000000..23b8c19f9e --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/stable_diffusion/get_vae_scale.rs @@ -0,0 +1,13 @@ +use crate::ml::model_file::StableDiffusionVersion; + +pub fn get_vae_scale(sd_version: StableDiffusionVersion) -> f64 { + match sd_version { + StableDiffusionVersion::V1_5 + | StableDiffusionVersion::V1_5Inpaint + | StableDiffusionVersion::V2_1 + | StableDiffusionVersion::V2Inpaint + | StableDiffusionVersion::XlInpaint + | StableDiffusionVersion::Xl => 0.18215, + StableDiffusionVersion::Turbo => 0.13025, + } +} diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/mod.rs b/crates/desktop/src-tauri/src/ml/stable_diffusion/mod.rs new file mode 100644 index 0000000000..8e9407a18e --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/stable_diffusion/mod.rs @@ -0,0 +1 @@ +pub mod get_vae_scale; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/state/app_config.rs b/crates/desktop/src-tauri/src/state/app_config.rs index 894c502955..3c5f5e0a7c 100644 --- a/crates/desktop/src-tauri/src/state/app_config.rs +++ b/crates/desktop/src-tauri/src/state/app_config.rs @@ -40,10 +40,9 @@ pub struct AppConfig { impl AppConfig { pub fn init() -> anyhow::Result { - + println!("Configuration filename: {}", CONFIG_FILENAME); let yaml_configs = load_yaml_configs(); - println!("Configured filename {}", CONFIG_FILENAME); println!("Configs: {:?}", yaml_configs); let device = Device::new_cuda(0) @@ -83,6 +82,7 @@ fn load_yaml_configs() -> YamlConfig { .ok() .flatten(); if let Some(path) = maybe_path { + let path = path.canonicalize().unwrap_or(path); println!("Attempting to load configs from: {:?}", &path); if let Ok(config) = YamlConfig::read_from_file(path) { return config;