From 1bc275293bd32972238d64db4ce34af80a5a0eb5 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 3 Aug 2025 15:45:09 -0400 Subject: [PATCH] masking and inpainting works --- .../enqueue_image_inpaint_command.rs | 4 +- .../handle_flux_pro_1_inpaint_artcraft.rs | 5 +- crates/lib/files/images/src/lib.rs | 1 + .../lib/files/images/src/mask_images/mod.rs | 1 + .../normalize_image_bytes_to_flux_mask.rs | 67 +++++++++++++++++++ .../libs/components/modal/src/lib/modal.tsx | 2 - 6 files changed, 74 insertions(+), 6 deletions(-) create mode 100644 crates/lib/files/images/src/mask_images/mod.rs create mode 100644 crates/lib/files/images/src/mask_images/normalize_image_bytes_to_flux_mask.rs diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs index 74ae149cf5..78d5e749a8 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs @@ -131,8 +131,8 @@ pub async fn enqueue_image_inpaint_command( sora_task_queue: State<'_, SoraTaskQueue>, ) -> ResponseOrErrorType { - info!("enqueue_image_inpaint_command called; full request: {:?}", - &request); + info!("enqueue_image_inpaint_command called; model: {:?}, image_media_token: {:?}, prompt: {:?}", + &request.model, &request.image_media_token, &request.prompt); let result = handle_request( &request, diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/flux_pro_1_inpaint/handle_flux_pro_1_inpaint_artcraft.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/flux_pro_1_inpaint/handle_flux_pro_1_inpaint_artcraft.rs index 47f43d52ee..fc16c853c9 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/flux_pro_1_inpaint/handle_flux_pro_1_inpaint_artcraft.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/flux_pro_1_inpaint/handle_flux_pro_1_inpaint_artcraft.rs @@ -47,6 +47,7 @@ use storyteller_client::media_files::get_media_file::get_media_file; use storyteller_client::media_files::upload_image_media_file_from_bytes::upload_image_media_file_from_bytes; use storyteller_client::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file; use tauri::AppHandle; +use images::mask_images::normalize_image_bytes_to_flux_mask::normalize_image_bytes_to_flux_mask; use tokens::tokens::media_files::MediaFileToken; pub async fn handle_flux_pro_1_inpaint_artcraft( @@ -156,7 +157,7 @@ async fn get_mask( let image_bytes = request.mask_image_raw_bytes.as_ref() .ok_or(InternalImageInpaintError::NoMaskImageSpecified)?; - let png_with_dims = image_bytes_to_png_bytes_with_dimensions(image_bytes) + let image_bytes = normalize_image_bytes_to_flux_mask(image_bytes) .map_err(|err| { error!("Failed to convert image bytes to png: {:?}", err); InternalImageInpaintError::CouldNotEncodeMask @@ -167,7 +168,7 @@ async fn get_mask( let result = upload_image_media_file_from_bytes( &app_env_configs.storyteller_host, Some(&storyteller_creds), - png_with_dims.png_bytes + image_bytes.0 ).await .map_err(|err| { error!("Failed to upload image media file: {:?}", err); diff --git a/crates/lib/files/images/src/lib.rs b/crates/lib/files/images/src/lib.rs index 97a83f180d..c65665cf34 100644 --- a/crates/lib/files/images/src/lib.rs +++ b/crates/lib/files/images/src/lib.rs @@ -26,6 +26,7 @@ pub use image; pub mod encoding; pub mod error; pub mod image_info; +pub mod mask_images; pub mod resize_image_file_preserving_aspect; pub mod resize_preserving_aspect; diff --git a/crates/lib/files/images/src/mask_images/mod.rs b/crates/lib/files/images/src/mask_images/mod.rs new file mode 100644 index 0000000000..b2b6d6f17f --- /dev/null +++ b/crates/lib/files/images/src/mask_images/mod.rs @@ -0,0 +1 @@ +pub mod normalize_image_bytes_to_flux_mask; \ No newline at end of file diff --git a/crates/lib/files/images/src/mask_images/normalize_image_bytes_to_flux_mask.rs b/crates/lib/files/images/src/mask_images/normalize_image_bytes_to_flux_mask.rs new file mode 100644 index 0000000000..edbcee6a8d --- /dev/null +++ b/crates/lib/files/images/src/mask_images/normalize_image_bytes_to_flux_mask.rs @@ -0,0 +1,67 @@ +use crate::error::ImagesError; +use image::{GenericImage, GenericImageView, ImageFormat, ImageReader, Rgba}; +use std::io::Cursor; + +pub struct PngBytes(pub Vec); + +/// Normalize arbitrary image bytes to a PNG format image. +/// Make sure the bitmap is in the correct mask format. +pub fn normalize_image_bytes_to_flux_mask( + arbitrary_image_bytes: &[u8], +) -> Result { + + let reader = ImageReader::new(Cursor::new(arbitrary_image_bytes)); + + let mut image = reader + .with_guessed_format()? // NB: Can raise IoError. + .decode()?; + + let (width, height) = image.dimensions(); + + //let mut corrected_mask: ImageBuffer, Vec> = ImageBuffer::new(width, height); + + // Iterate through each pixel + for y in 0..height { + for x in 0..width { + let rgba_pixel = image.get_pixel(x, y); + + if rgba_pixel[3] == 0 { + // Alpha pixel is transparent + // Set to black. Mask ignore region. + image.put_pixel(x, y, Rgba([0, 0, 0, 255])); + } else { + // Set to white. Mask interest region. + image.put_pixel(x, y, Rgba([255, 255, 255, 255])); + } + + /*match pixel.channels().get(4) { + Some(0) => { + // Alpha pixel is transparent + // Set to black. Mask ignore region. + pixel = Rgba([0, 0, 0, 255]); + } + _ => { + // Set to white. Mask interest region. + pixel = Rgba([255, 255, 255, 255]); + }, + } + + let rgb_pixel = pixel.to_rgba(); // Convert to Rgba + + // Check if the pixel is pure black (R=0, G=0, B=0) + if rgb_pixel[0] == 0 && rgb_pixel[1] == 0 && rgb_pixel[2] == 0 { + // Set the pixel to white (R=255, G=255, B=255) + corrected_mask.put_pixel(x, y, Rgba([255, 255, 255, 255])); + } else { + // Keep the original color for non-black pixels + corrected_mask.put_pixel(x, y, rgb_pixel); + }*/ + } + } + + + let mut output_bytes: Vec = Vec::new(); + image.write_to(&mut Cursor::new(&mut output_bytes), ImageFormat::Png)?; + + Ok(PngBytes(output_bytes)) +} diff --git a/frontend/libs/components/modal/src/lib/modal.tsx b/frontend/libs/components/modal/src/lib/modal.tsx index 91c83390cd..4eaf6cec4a 100644 --- a/frontend/libs/components/modal/src/lib/modal.tsx +++ b/frontend/libs/components/modal/src/lib/modal.tsx @@ -895,8 +895,6 @@ export const Modal = ({ > {/* Animated Modal Content */} {transitions((styles, item) => { - // Debug logging - remove this after testing - console.log("Modal animation styles:", styles); return item ? (