diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/utils/mod.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/utils/mod.rs new file mode 100644 index 0000000000..1734914faa --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/utils/mod.rs @@ -0,0 +1 @@ +pub mod parse_semantic_media_files; diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/utils/parse_semantic_media_files.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/utils/parse_semantic_media_files.rs new file mode 100644 index 0000000000..7b604958c3 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/utils/parse_semantic_media_files.rs @@ -0,0 +1,123 @@ +use log::info; + +use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet; +use artcraft_client::utils::api_host::ApiHost; +use tokens::tokens::media_files::MediaFileToken; + +use crate::core::commands::enqueue::generate_error::{BadInputReason, GenerateError}; +use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest; +use crate::core::utils::upload_bytes_to_media_file::upload_image_bytes_as_media_file::upload_image_bytes_as_media_file; +use crate::core::utils::upload_bytes_to_media_file::upload_mask_image_bytes_as_media_file::upload_mask_image_bytes_as_media_file; + +/// Resolved semantic media file tokens for canvas, scene, and inpainting mask. +/// +/// Each field is populated either from an already-uploaded token on the request +/// or by uploading raw bytes on the fly. +pub struct SemanticMediaFiles { + pub canvas_image_media_token: Option, + pub scene_image_media_token: Option, + pub inpainting_mask_image_media_token: Option, +} + +/// Resolve semantic media files from the request. +/// +/// For each of (canvas, scene, mask): the request may supply a pre-uploaded +/// token XOR raw bytes. If both are supplied, returns an error. If raw bytes +/// are supplied, uploads them and returns the resulting token. +pub async fn parse_semantic_media_files( + request: &TauriGenerateImageRequest, + creds: &StorytellerCredentialSet, + api_host: &ApiHost, +) -> Result { + // Canvas image + let canvas_image_media_token = resolve_image_field( + "canvas_image", + request.canvas_image_media_token.as_ref(), + request.canvas_image_raw_bytes.as_deref(), + creds, + api_host, + ).await?; + + // Scene image + let scene_image_media_token = resolve_image_field( + "scene_image", + request.scene_image_media_token.as_ref(), + request.scene_image_raw_bytes.as_deref(), + creds, + api_host, + ).await?; + + // Inpainting mask + let inpainting_mask_image_media_token = resolve_mask_field( + request.inpainting_mask_image_media_token.as_ref(), + request.inpainting_mask_image_raw_bytes.as_deref(), + creds, + api_host, + ).await?; + + Ok(SemanticMediaFiles { + canvas_image_media_token, + scene_image_media_token, + inpainting_mask_image_media_token, + }) +} + +// ── Helpers ── + +async fn resolve_image_field( + field_name: &str, + maybe_token: Option<&MediaFileToken>, + maybe_raw_bytes: Option<&[u8]>, + creds: &StorytellerCredentialSet, + api_host: &ApiHost, +) -> Result, GenerateError> { + if maybe_token.is_some() && maybe_raw_bytes.is_some() { + return Err(GenerateError::BadInput(BadInputReason::WrongImageArguments(format!( + "Cannot supply both {}_media_token and {}_raw_bytes", + field_name, field_name + )))); + } + + if let Some(token) = maybe_token { + return Ok(Some(token.clone())); + } + + if let Some(raw_bytes) = maybe_raw_bytes { + info!("Uploading {} raw bytes ({} bytes)...", field_name, raw_bytes.len()); + let token = upload_image_bytes_as_media_file( + creds, + api_host, + raw_bytes.to_vec(), + ).await?; + return Ok(Some(token)); + } + + Ok(None) +} + +async fn resolve_mask_field( + maybe_token: Option<&MediaFileToken>, + maybe_raw_bytes: Option<&[u8]>, + creds: &StorytellerCredentialSet, + api_host: &ApiHost, +) -> Result, GenerateError> { + if maybe_token.is_some() && maybe_raw_bytes.is_some() { + return Err(GenerateError::BadInput(BadInputReason::BothImageMaskMediaTokenAndBytesSupplied)); + } + + if let Some(token) = maybe_token { + return Ok(Some(token.clone())); + } + + if let Some(raw_bytes) = maybe_raw_bytes { + info!("Uploading inpainting mask raw bytes ({} bytes)...", raw_bytes.len()); + let token = upload_mask_image_bytes_as_media_file( + creds, + api_host, + raw_bytes, + ).await?; + return Ok(Some(token)); + } + + Ok(None) +}