diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index 1ce106b19e..d5c100d29d 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs @@ -6,6 +6,7 @@ use crate::model_config::ModelConfig; use anyhow::Result; use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; use image::ImageReader; +use log::error; use tauri::State; const PROMPT: &str = "A beautiful landscape with mountains and a lake"; @@ -31,7 +32,13 @@ pub fn infer_image(image: &str, model_config: State, model_cache: S image.save(&image_path) .map_err(|err| format!("Failed to save image: {}", err))?; - do_infer_image(&image_path, &model_config, &model_cache) + let result = do_infer_image(&image_path, &model_config, &model_cache); + + if let Err(err) = result.as_deref() { + error!("There was an error: {:?}", err); + } + + result } fn do_infer_image(image_path: &PathBuf, config: &ModelConfig, model_cache: &ModelCache) -> Result { diff --git a/crates/desktop/src-tauri/src/endpoints/sd.rs b/crates/desktop/src-tauri/src/endpoints/sd.rs index 285ecdf2d2..66885434f9 100644 --- a/crates/desktop/src-tauri/src/endpoints/sd.rs +++ b/crates/desktop/src-tauri/src/endpoints/sd.rs @@ -14,6 +14,7 @@ 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 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; @@ -140,8 +141,15 @@ pub fn run(args: Args<'_>) -> Result<()> { // TODO: Just to trip the conditionals that force mask/inpaint generation // TODO: Set this to turbo to turn off inpaint, anything else will turn it on. - const HACK_INPAINT_SD_VERSION : StableDiffusionVersion = StableDiffusionVersion::Turbo; + // TODO: THIS MIGHT NOT BE POSSIBLE WITH TURBO: https://huggingface.co/stabilityai/sdxl-turbo/discussions/7 + // TODO/SEE ALSO: https://www.reddit.com/r/StableDiffusion/comments/198cl08/major_inpainting_issues_with_sdxl_xl_turbo/ + // TODO/SEE ALSO: Alternate weights/issue: https://github.com/huggingface/diffusers/issues/6529 + // TODO/SEE ALSO: https://huggingface.co/spaces/OzzyGT/diffusers-fast-inpaint (THIS LOOKS AWFUL THOUGH -- let's not use this) + // + // TODO: Referencing this code: https://github.com/placrosse/candle/blob/43017539ab4f9ccb43015b456136b704ebf693e0/candle-examples/examples/stable-diffusion/main.rs#L491 + const HACK_INPAINT_SD_VERSION : StableDiffusionVersion = StableDiffusionVersion::XlInpaint; + // TODO: This is inpainting. It doesn't work yet. let (mask_latents, mask, mask_4) = create_inpainting_tensors( HACK_INPAINT_SD_VERSION, // TODO: This is a hack Some(args.image_path.to_path_buf()), // TODO: Mask needs hardcoding @@ -156,6 +164,7 @@ pub fn run(args: Args<'_>) -> Result<()> { println!("Generating initial noise..."); let timesteps = scheduler.timesteps().to_vec(); + // TODO: This shape may need to change for inpainting. let latents = Tensor::randn( 0f32, 1f32, @@ -186,14 +195,20 @@ pub fn run(args: Args<'_>) -> Result<()> { let latent_model_input = match HACK_INPAINT_SD_VERSION { StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint - | StableDiffusionVersion::V1_5Inpaint => Tensor::cat( - &[ - &latent_model_input, - mask.as_ref().unwrap(), - mask_latents.as_ref().unwrap(), - ], - 1, - )?, + | StableDiffusionVersion::V1_5Inpaint => { + info!("Concatenating input shape: {:?}", latent_model_input.shape()); + info!("Mask shape: {:?}", mask.as_ref().unwrap().shape()); + info!("Mask latents shape: {:?}", mask_latents.as_ref().unwrap().shape()); + info!("IF THIS FAILS, REVERT THE `HACK_INPAINT_SD_VERSION` 'flag'"); + Tensor::cat( + &[ + &latent_model_input, + mask.as_ref().unwrap(), + mask_latents.as_ref().unwrap(), + ], + 1, + )? + }, _ => latent_model_input, } .to_device(&args.model_configs.device)?;