From 588d30d2cea49b33ab407db75bebf5b4dcbf77f3 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 May 2026 01:44:41 -0400 Subject: [PATCH] wip --- .../generate_image/generate_image_command.rs | 74 ++++++++++++++-- .../handle_artcraft_via_omni_endpoint.rs | 20 +++-- .../providers/artcraft/model_mapping.rs | 88 +++++++++++++++---- .../tauri_generate_image_request.rs | 4 +- 4 files changed, 151 insertions(+), 35 deletions(-) diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/generate_image_command.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/generate_image_command.rs index b514035c48..96976339e4 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/generate_image_command.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/generate_image_command.rs @@ -15,6 +15,7 @@ use crate::services::storyteller::state::storyteller_credential_manager::Storyte use enums::common::generation_provider::GenerationProvider; use log::{error, info}; use tauri::{AppHandle, State}; +use crate::core::commands::enqueue::common::notify_frontend_of_errors::notify_frontend_of_errors; #[tauri::command] pub async fn generate_image_command( @@ -50,10 +51,21 @@ pub async fn generate_image_command( &storyteller_creds_manager, ).await; - let success = match result { - Ok(success) => success, - Err(err) => return map_error_to_response(err), - }; + match result { + Ok(success) => handle_success_behavior(&app, &task_database, &request, success).await, + Err(err) => handle_error_behavior(&app, err).await, + } +} + +// ── Result mapping ── + +async fn handle_success_behavior( + app: &AppHandle, + task_database: &TaskDatabase, + request: &TauriGenerateImageRequest, + success: TaskEnqueueSuccess, +) -> Response { + // Insert into task database let db_result = success @@ -69,6 +81,19 @@ pub async fn generate_image_command( error!("Failed to create task in database: {:?}", err); } + /* + let usage_type = if is_image_to_image { + ArtcraftUsageType::ImageToResult + } else { + ArtcraftUsageType::TextToResult + }; + + if let Err(err) = artcraft_usage_tracker.record_image_generation(num_images, usage_type, ArtcraftUsagePage::ImagePage) { + // NB: Fail open. + warn!("Failed to report usage: {:?}", err); + } + */ + let event = GenerationEnqueueSuccessEvent { action: success.to_frontend_event_action(), service: success.to_frontend_event_service(), @@ -84,14 +109,14 @@ pub async fn generate_image_command( Ok(TauriGenerateImageResponse {}.into()) } -// ── Result mapping ── - - -fn map_error_to_response( +async fn handle_error_behavior( + app: &AppHandle, err: GenerateError, ) -> Response { error!("generate_image_command error: {:?}", err); + notify_frontend_of_errors(&app, &err).await; + let error_type = match &err { GenerateError::BadInput(_) => TauriGenerateImageErrorType::BadInput, GenerateError::MissingCredentials(_) => TauriGenerateImageErrorType::NeedsStorytellerCredentials, @@ -100,6 +125,39 @@ fn map_error_to_response( _ => TauriGenerateImageErrorType::ServerError, }; + // TODO: + /* + let mut status = CommandErrorStatus::ServerError; + let mut error_type = EnqueueTextToImageErrorType::ServerError; + let mut error_message = "A server error occurred. Please try again. If it continues, please tell our staff about the problem."; + + match err { + GenerateError::BadInput(BadInputReason::NoModelSpecified) => { + status = CommandErrorStatus::BadRequest; + error_type = EnqueueTextToImageErrorType::ModelNotSpecified; + error_message = "No model specified for image generation"; + } + GenerateError::NoProviderAvailable => { + status = CommandErrorStatus::ServerError; + error_type = EnqueueTextToImageErrorType::NoProviderAvailable; + error_message = "No configured provider available for image generation"; + } + GenerateError::MissingCredentials(MissingCredentialsReason::NeedsFalApiKey) => { + status = CommandErrorStatus::Unauthorized; + error_type = EnqueueTextToImageErrorType::NeedsFalApiKey; + error_message = "You need to set a FAL api key"; + }, + _ => {}, // Fall-through + } + + Err(CommandErrorResponseWrapper { + status, + error_message: Some(error_message.to_string()), + error_type: Some(error_type), + error_details: None, + }) + */ + Err(CommandErrorResponseWrapper { status: CommandErrorStatus::ServerError, error_message: Some(format!("{:?}", err)), diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_omni_endpoint.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_omni_endpoint.rs index c02fb1a586..1071bbf99d 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_omni_endpoint.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_omni_endpoint.rs @@ -3,11 +3,11 @@ use artcraft_client::endpoints::omni_gen::generate::image::omni_gen_image::omni_ use enums::common::generation_provider::GenerationProvider; use enums::tauri::tasks::task_type::TaskType; use log::{error, info}; - +use uuid_utils::uuid::generate_random_uuid; use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason}; use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; use crate::core::commands::generate::generate_image::providers::artcraft::model_mapping::{ - map_to_generation_model, map_to_omni_image_model, + map_to_generation_model, map_to_omni_aspect_ratio, map_to_omni_image_model, map_to_omni_resolution, }; use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; @@ -18,20 +18,22 @@ pub async fn handle_artcraft_via_omni_endpoint( app_env_configs: &AppEnvConfigs, storyteller_creds_manager: &StorytellerCredentialManager, ) -> Result { - let model = request.model.ok_or(GenerateError::no_model_specified())?; + let tauri_model = request.model.ok_or(GenerateError::no_model_specified())?; - let omni_model = map_to_omni_image_model(model) + let omni_model = map_to_omni_image_model(tauri_model) .ok_or(GenerateError::NotYetImplemented( - format!("Model {:?} is not supported via the omni endpoint", model), + format!("Model {:?} is not supported via the omni endpoint", tauri_model), ))?; - let generation_model = map_to_generation_model(model); + let generation_model = map_to_generation_model(tauri_model); let creds = match storyteller_creds_manager.get_credentials()? { Some(creds) => creds, None => return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsStorytellerCredentials)), }; + let uuid_idempotency_token = generate_random_uuid(); + // Build image_media_tokens: prepend canvas image if present. let mut image_media_tokens = request.image_media_tokens.clone().unwrap_or_default(); if let Some(canvas_token) = &request.canvas_image_media_token { @@ -40,12 +42,12 @@ pub async fn handle_artcraft_via_omni_endpoint( let image_media_tokens = if image_media_tokens.is_empty() { None } else { Some(image_media_tokens) }; let omni_request = OmniGenImageCostAndGenerateRequest { - idempotency_token: None, + idempotency_token: Some(uuid_idempotency_token), model: Some(omni_model), prompt: request.prompt.clone(), image_media_tokens, - resolution: None, - aspect_ratio: None, + resolution: request.resolution, + aspect_ratio: request.aspect_ratio, quality: request.quality, image_batch_count: request.batch_size.map(|n| n as u16), adjust_horizontal_angle: request.adjust_horizontal_angle, diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/model_mapping.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/model_mapping.rs index fa6fee7e07..1e26f8e5da 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/model_mapping.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/model_mapping.rs @@ -1,5 +1,9 @@ +use artcraft_router::api::common_aspect_ratio::CommonAspectRatio as RouterCommonAspectRatio; use artcraft_router::api::common_image_model::CommonImageModel as RouterCommonImageModel; +use artcraft_router::api::common_resolution::CommonResolution as RouterCommonResolution; +use enums::common::generation::common_aspect_ratio::CommonAspectRatio as OmniCommonAspectRatio; use enums::common::generation::common_image_model::CommonImageModel as OmniCommonImageModel; +use enums::common::generation::common_resolution::CommonResolution as OmniCommonResolution; use crate::core::commands::generate::generate_image::tauri_image_model::TauriImageModel; use crate::core::events::generation_events::common::GenerationModel; @@ -10,6 +14,7 @@ pub fn map_to_omni_image_model(model: TauriImageModel) -> Option Some(OmniCommonImageModel::Flux1Dev), TauriImageModel::Flux1Schnell => Some(OmniCommonImageModel::Flux1Schnell), + TauriImageModel::FluxPro1 => Some(OmniCommonImageModel::FluxPro11), // TODO: Might be a slight mismatch TauriImageModel::FluxPro11 => Some(OmniCommonImageModel::FluxPro11), TauriImageModel::FluxPro11Ultra => Some(OmniCommonImageModel::FluxPro11Ultra), TauriImageModel::GptImage1 => Some(OmniCommonImageModel::GptImage1), @@ -22,7 +27,14 @@ pub fn map_to_omni_image_model(model: TauriImageModel) -> Option Some(OmniCommonImageModel::Seedream4), TauriImageModel::Seedream4p5 => Some(OmniCommonImageModel::Seedream4p5), TauriImageModel::Seedream5Lite => Some(OmniCommonImageModel::Seedream5Lite), - _ => None, + // Not accounted for yet + TauriImageModel::GrokImage => None, + TauriImageModel::Recraft3 => None, + TauriImageModel::Midjourney => None, + TauriImageModel::FluxProKontextMax => None, + TauriImageModel::QwenEdit2511Angles => None, + TauriImageModel::Flux2LoraAngles => None, + TauriImageModel::FluxDevJuggernaut => None, } } @@ -30,23 +42,67 @@ pub fn map_to_omni_image_model(model: TauriImageModel) -> Option Option { match model { - TauriImageModel::Flux1Dev => Some(RouterCommonImageModel::Flux1Dev), - TauriImageModel::Flux1Schnell => Some(RouterCommonImageModel::Flux1Schnell), - TauriImageModel::FluxPro11 => Some(RouterCommonImageModel::FluxPro11), - TauriImageModel::FluxPro11Ultra => Some(RouterCommonImageModel::FluxPro11Ultra), - TauriImageModel::GptImage1 => Some(RouterCommonImageModel::GptImage1), - TauriImageModel::GptImage1p5 => Some(RouterCommonImageModel::GptImage1p5), - TauriImageModel::GptImage2 => Some(RouterCommonImageModel::GptImage2), - TauriImageModel::NanoBanana => Some(RouterCommonImageModel::NanoBanana), - TauriImageModel::NanoBanana2 => Some(RouterCommonImageModel::NanoBanana2), - TauriImageModel::NanoBananaPro => Some(RouterCommonImageModel::NanoBananaPro), - TauriImageModel::Gemini25Flash => Some(RouterCommonImageModel::NanoBanana), - TauriImageModel::Seedream4 => Some(RouterCommonImageModel::Seedream4), - TauriImageModel::Seedream4p5 => Some(RouterCommonImageModel::Seedream4p5), - TauriImageModel::Seedream5Lite => Some(RouterCommonImageModel::Seedream5Lite), + TauriImageModel::Flux1Dev => Some(RouterCommonImageModel::Flux1Dev), // Text-to-Image + TauriImageModel::Flux1Schnell => Some(RouterCommonImageModel::Flux1Schnell), // Text-to-Image + TauriImageModel::FluxPro1 => Some(RouterCommonImageModel::FluxPro11), // TODO: Might be a slight mismatch + TauriImageModel::FluxPro11 => Some(RouterCommonImageModel::FluxPro11), // Text-to-Image + TauriImageModel::FluxPro11Ultra => Some(RouterCommonImageModel::FluxPro11Ultra), // Text-to-Image + TauriImageModel::GptImage1 => Some(RouterCommonImageModel::GptImage1), // Text-to-Image + TauriImageModel::GptImage1p5 => Some(RouterCommonImageModel::GptImage1p5), // Text-to-Image + TauriImageModel::GptImage2 => Some(RouterCommonImageModel::GptImage2), // Text-to-Image + TauriImageModel::NanoBanana => Some(RouterCommonImageModel::NanoBanana), // Text-to-Image + TauriImageModel::NanoBanana2 => Some(RouterCommonImageModel::NanoBanana2), // Text-to-Image + TauriImageModel::NanoBananaPro => Some(RouterCommonImageModel::NanoBananaPro), // Text-to-Image + TauriImageModel::Gemini25Flash => Some(RouterCommonImageModel::NanoBanana), // Text-to-Image + TauriImageModel::Seedream4 => Some(RouterCommonImageModel::Seedream4), // Text-to-Image + TauriImageModel::Seedream4p5 => Some(RouterCommonImageModel::Seedream4p5), // Text-to-Image + TauriImageModel::Seedream5Lite => Some(RouterCommonImageModel::Seedream5Lite), // Text-to-Image TauriImageModel::QwenEdit2511Angles => Some(RouterCommonImageModel::QwenEdit2511Angles), TauriImageModel::Flux2LoraAngles => Some(RouterCommonImageModel::Flux2LoraAngles), - _ => None, + // Not accounted for yet + TauriImageModel::GrokImage => None, + TauriImageModel::Recraft3 => None, + TauriImageModel::Midjourney => None, + TauriImageModel::FluxProKontextMax => None, + TauriImageModel::FluxDevJuggernaut => None, + } +} + +/// Map from the router's CommonAspectRatio to the omni endpoint's CommonAspectRatio. +pub fn map_to_omni_aspect_ratio(ratio: RouterCommonAspectRatio) -> OmniCommonAspectRatio { + match ratio { + RouterCommonAspectRatio::Auto => OmniCommonAspectRatio::Auto, + RouterCommonAspectRatio::Square => OmniCommonAspectRatio::Square, + RouterCommonAspectRatio::WideThreeByTwo => OmniCommonAspectRatio::WideThreeByTwo, + RouterCommonAspectRatio::WideFourByThree => OmniCommonAspectRatio::WideFourByThree, + RouterCommonAspectRatio::WideFiveByFour => OmniCommonAspectRatio::WideFiveByFour, + RouterCommonAspectRatio::WideSixteenByNine => OmniCommonAspectRatio::WideSixteenByNine, + RouterCommonAspectRatio::WideTwentyOneByNine => OmniCommonAspectRatio::WideTwentyOneByNine, + RouterCommonAspectRatio::TallTwoByThree => OmniCommonAspectRatio::TallTwoByThree, + RouterCommonAspectRatio::TallThreeByFour => OmniCommonAspectRatio::TallThreeByFour, + RouterCommonAspectRatio::TallFourByFive => OmniCommonAspectRatio::TallFourByFive, + RouterCommonAspectRatio::TallNineBySixteen => OmniCommonAspectRatio::TallNineBySixteen, + RouterCommonAspectRatio::TallNineByTwentyOne => OmniCommonAspectRatio::TallNineByTwentyOne, + RouterCommonAspectRatio::Wide => OmniCommonAspectRatio::Wide, + RouterCommonAspectRatio::Tall => OmniCommonAspectRatio::Tall, + RouterCommonAspectRatio::Auto2k => OmniCommonAspectRatio::Auto2k, + RouterCommonAspectRatio::Auto3k => OmniCommonAspectRatio::Auto3k, + RouterCommonAspectRatio::Auto4k => OmniCommonAspectRatio::Auto4k, + RouterCommonAspectRatio::SquareHd => OmniCommonAspectRatio::SquareHd, + } +} + +/// Map from the router's CommonResolution to the omni endpoint's CommonResolution. +pub fn map_to_omni_resolution(res: RouterCommonResolution) -> OmniCommonResolution { + match res { + RouterCommonResolution::OneK => OmniCommonResolution::OneK, + RouterCommonResolution::TwoK => OmniCommonResolution::TwoK, + RouterCommonResolution::ThreeK => OmniCommonResolution::ThreeK, + RouterCommonResolution::FourK => OmniCommonResolution::FourK, + RouterCommonResolution::HalfK => OmniCommonResolution::HalfK, + RouterCommonResolution::FourEightyP => OmniCommonResolution::FourEightyP, + RouterCommonResolution::SevenTwentyP => OmniCommonResolution::SevenTwentyP, + RouterCommonResolution::TenEightyP => OmniCommonResolution::TenEightyP, } } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/tauri_generate_image_request.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/tauri_generate_image_request.rs index 759eb4274e..1a4efafaea 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/tauri_generate_image_request.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/tauri_generate_image_request.rs @@ -1,5 +1,5 @@ -use artcraft_router::api::common_aspect_ratio::CommonAspectRatio; -use artcraft_router::api::common_resolution::CommonResolution; +use enums::common::generation::common_aspect_ratio::CommonAspectRatio; +use enums::common::generation::common_resolution::CommonResolution; use enums::common::generation::common_quality::CommonQuality; use enums::common::generation_provider::GenerationProvider; use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;