From 13896d9306167a1c8f14b558fc4e26532ee79944 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 May 2026 01:22:58 -0400 Subject: [PATCH] cleanup --- .../generate_image/generate_image_command.rs | 50 +++++----- .../providers/artcraft/handle_artcraft.rs | 40 ++++++++ .../artcraft/handle_artcraft_via_legacy.rs | 95 +++++++++++++++++++ .../handle_artcraft_via_omni_endpoint.rs | 68 ++----------- .../generate_image/providers/artcraft/mod.rs | 7 +- .../providers/artcraft/model_mapping.rs | 79 +++++++++++++++ 6 files changed, 249 insertions(+), 90 deletions(-) create mode 100644 crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft.rs create mode 100644 crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_legacy.rs create mode 100644 crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/model_mapping.rs 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 4a6a962654..b514035c48 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 @@ -1,6 +1,6 @@ use crate::core::commands::enqueue::generate_error::GenerateError; use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; -use crate::core::commands::generate::generate_image::providers::artcraft::handle_artcraft_via_omni_endpoint::handle_artcraft_via_omni_endpoint; +use crate::core::commands::generate::generate_image::providers::artcraft::handle_artcraft; use crate::core::commands::generate::generate_image::tauri_generate_image_request::{ TauriGenerateImageErrorType, TauriGenerateImageRequest, TauriGenerateImageResponse, }; @@ -44,55 +44,49 @@ pub async fn generate_image_command( }); } - let result = handle_artcraft_via_omni_endpoint( + let result = handle_artcraft( &request, &app_env_configs, &storyteller_creds_manager, ).await; - match result { - Ok(success) => { - // Insert into task database - let db_result = success - .insert_into_task_database_with_frontend_payload( - &task_database, - request.frontend_caller, - request.frontend_subscriber_id.as_deref(), - request.frontend_subscriber_payload.as_deref(), - ) - .await; + let success = match result { + Ok(success) => success, + Err(err) => return map_error_to_response(err), + }; - if let Err(err) = db_result { - error!("Failed to create task in database: {:?}", err); - } + // Insert into task database + let db_result = success + .insert_into_task_database_with_frontend_payload( + &task_database, + request.frontend_caller, + request.frontend_subscriber_id.as_deref(), + request.frontend_subscriber_payload.as_deref(), + ) + .await; - map_success_to_response(success, &app) - } - Err(err) => map_error_to_response(err), + if let Err(err) = db_result { + error!("Failed to create task in database: {:?}", err); } -} -// ── Result mapping ── - -fn map_success_to_response( - success: TaskEnqueueSuccess, - app: &AppHandle, -) -> Response { let event = GenerationEnqueueSuccessEvent { action: success.to_frontend_event_action(), service: success.to_frontend_event_service(), model: success.model, }; - if let Err(err) = event.send(app) { + if let Err(err) = event.send(&app) { error!("Failed to emit event: {:?}", err); } - CreditsBalanceChangedEvent{}.send_infallible(app); + CreditsBalanceChangedEvent{}.send_infallible(&app); Ok(TauriGenerateImageResponse {}.into()) } +// ── Result mapping ── + + fn map_error_to_response( err: GenerateError, ) -> Response { diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft.rs new file mode 100644 index 0000000000..0e10e5765d --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft.rs @@ -0,0 +1,40 @@ +use log::info; + +use crate::core::commands::enqueue::generate_error::GenerateError; +use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; +use crate::core::commands::generate::generate_image::providers::artcraft::handle_artcraft_via_legacy::handle_artcraft_via_legacy; +use crate::core::commands::generate::generate_image::providers::artcraft::handle_artcraft_via_omni_endpoint::handle_artcraft_via_omni_endpoint; +use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; + +/// Dispatch an image generation request to Artcraft. +/// +/// Tries the omni endpoint first. If the model isn't supported there, +/// falls back to the legacy artcraft_router path. +pub async fn handle_artcraft( + request: &TauriGenerateImageRequest, + app_env_configs: &AppEnvConfigs, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> Result { + // Try the omni endpoint first. + let omni_result = handle_artcraft_via_omni_endpoint( + request, + app_env_configs, + storyteller_creds_manager, + ).await; + + match omni_result { + Ok(success) => Ok(success), + Err(GenerateError::NotYetImplemented(_)) => { + // Model not supported by omni endpoint — fall back to legacy. + info!("Model not supported by omni endpoint, falling back to legacy artcraft_router path."); + handle_artcraft_via_legacy( + request, + app_env_configs, + storyteller_creds_manager, + ).await + } + Err(err) => Err(err), + } +} diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_legacy.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_legacy.rs new file mode 100644 index 0000000000..d45c8f2bba --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_legacy.rs @@ -0,0 +1,95 @@ +use artcraft_router::api::common_image_model::CommonImageModel; +use artcraft_router::api::image_list_ref::ImageListRef; +use artcraft_router::api::provider::Provider; +use artcraft_router::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy; +use artcraft_router::client::router_artcraft_client::RouterArtcraftClient; +use artcraft_router::client::router_client::RouterClient; +use artcraft_router::generate::generate_image::generate_image_request::GenerateImageRequest; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_type::TaskType; +use log::{error, info}; + +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_router_image_model, +}; +use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; + +/// Handle image generation via the legacy artcraft_router path. +/// +/// This is used for models not yet supported by the omni endpoint +/// (eg. Recraft, FluxProKontextMax, edit/inpaint models). +pub async fn handle_artcraft_via_legacy( + request: &TauriGenerateImageRequest, + app_env_configs: &AppEnvConfigs, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> Result { + let model = request.model.ok_or(GenerateError::no_model_specified())?; + + let router_model = map_to_router_image_model(model) + .ok_or(GenerateError::NotYetImplemented( + format!("Model {:?} is not supported via the legacy router path", model), + ))?; + + let generation_model = map_to_generation_model(model); + + let creds = match storyteller_creds_manager.get_credentials()? { + Some(creds) => creds, + None => return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsStorytellerCredentials)), + }; + + let client = RouterClient::Artcraft(RouterArtcraftClient::new( + app_env_configs.storyteller_host.clone(), + creds, + )); + + // Build image inputs: prepend canvas image if present. + let mut image_tokens = request.image_media_tokens.clone().unwrap_or_default(); + if let Some(canvas_token) = &request.canvas_image_media_token { + image_tokens.insert(0, canvas_token.clone()); + } + let image_inputs = if image_tokens.is_empty() { None } else { Some(ImageListRef::MediaFileTokens(image_tokens)) }; + + let router_request = GenerateImageRequest { + model: router_model, + provider: Provider::Artcraft, + prompt: request.prompt.clone(), + image_inputs, + resolution: None, // TODO: Convert from request.resolution + aspect_ratio: None, // TODO: Convert from request.aspect_ratio + quality: None, // TODO: Convert from request.quality + image_batch_count: request.batch_size.map(|n| n as u16), + request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade, + generation_mode_mismatch_strategy: None, + idempotency_token: None, + horizontal_angle: request.adjust_horizontal_angle, + vertical_angle: request.adjust_vertical_angle, + zoom: request.adjust_zoom, + }; + + let plan = router_request.build()?; + + info!("Legacy image generation plan: {:?}", plan); + + let response = plan.generate_image(&client).await.map_err(|err| { + error!("Legacy image generation failed: {:?}", err); + GenerateError::from(err) + })?; + + let job_id = response + .get_artcraft_payload() + .map(|p| p.inference_job_token.to_string()) + .ok_or(GenerateError::ResponseHadNoJobTokens)?; + + info!("Legacy image generation succeeded: job_id={}", job_id); + + Ok(TaskEnqueueSuccess { + task_type: TaskType::ImageGeneration, + model: Some(generation_model), + provider: GenerationProvider::Artcraft, + provider_job_id: Some(job_id), + }) +} 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 e32264a07c..c02fb1a586 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 @@ -1,15 +1,15 @@ use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; use artcraft_client::endpoints::omni_gen::generate::image::omni_gen_image::omni_gen_image_generate; -use enums::common::generation::common_image_model::CommonImageModel; use enums::common::generation_provider::GenerationProvider; use enums::tauri::tasks::task_type::TaskType; use log::{error, info}; 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, +}; use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest; -use crate::core::commands::generate::generate_image::tauri_image_model::TauriImageModel; -use crate::core::events::generation_events::common::GenerationModel; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; @@ -20,7 +20,7 @@ pub async fn handle_artcraft_via_omni_endpoint( ) -> Result { let model = request.model.ok_or(GenerateError::no_model_specified())?; - let omni_model = map_to_common_image_model(model) + let omni_model = map_to_omni_image_model(model) .ok_or(GenerateError::NotYetImplemented( format!("Model {:?} is not supported via the omni endpoint", model), ))?; @@ -32,14 +32,11 @@ pub async fn handle_artcraft_via_omni_endpoint( None => return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsStorytellerCredentials)), }; - // Build the omni gen request. + // Build image_media_tokens: prepend canvas image if present. let mut image_media_tokens = request.image_media_tokens.clone().unwrap_or_default(); - - // Prepend canvas image if present. if let Some(canvas_token) = &request.canvas_image_media_token { image_media_tokens.insert(0, canvas_token.clone()); } - let image_media_tokens = if image_media_tokens.is_empty() { None } else { Some(image_media_tokens) }; let omni_request = OmniGenImageCostAndGenerateRequest { @@ -47,8 +44,8 @@ pub async fn handle_artcraft_via_omni_endpoint( model: Some(omni_model), prompt: request.prompt.clone(), image_media_tokens, - resolution: None, // TODO: Convert from request.resolution - aspect_ratio: None, // TODO: Convert from request.aspect_ratio + resolution: None, + aspect_ratio: None, quality: request.quality, image_batch_count: request.batch_size.map(|n| n as u16), adjust_horizontal_angle: request.adjust_horizontal_angle, @@ -76,54 +73,3 @@ pub async fn handle_artcraft_via_omni_endpoint( provider_job_id: Some(response.inference_job_token.to_string()), }) } - -// ── Model mapping helpers ── - -fn map_to_common_image_model(model: TauriImageModel) -> Option { - match model { - TauriImageModel::Flux1Dev => Some(CommonImageModel::Flux1Dev), - TauriImageModel::Flux1Schnell => Some(CommonImageModel::Flux1Schnell), - TauriImageModel::FluxPro11 => Some(CommonImageModel::FluxPro11), - TauriImageModel::FluxPro11Ultra => Some(CommonImageModel::FluxPro11Ultra), - TauriImageModel::Gemini25Flash => Some(CommonImageModel::NanoBanana), // NB: Alias for Nano Banana - TauriImageModel::GptImage1 => Some(CommonImageModel::GptImage1), - TauriImageModel::GptImage1p5 => Some(CommonImageModel::GptImage1p5), - TauriImageModel::GptImage2 => Some(CommonImageModel::GptImage2), - TauriImageModel::NanoBanana => Some(CommonImageModel::NanoBanana), - TauriImageModel::NanoBanana2 => Some(CommonImageModel::NanoBanana2), - TauriImageModel::NanoBananaPro => Some(CommonImageModel::NanoBananaPro), - TauriImageModel::Seedream4 => Some(CommonImageModel::Seedream4), - TauriImageModel::Seedream4p5 => Some(CommonImageModel::Seedream4p5), - TauriImageModel::Seedream5Lite => Some(CommonImageModel::Seedream5Lite), - // Models not in CommonImageModel (Grok, Recraft, Midjourney, edit-only, inpaint-only) - _ => None, - } -} - -fn map_to_generation_model(model: TauriImageModel) -> GenerationModel { - match model { - TauriImageModel::Flux1Dev => GenerationModel::Flux1Dev, - TauriImageModel::Flux1Schnell => GenerationModel::Flux1Schnell, - TauriImageModel::Flux2LoraAngles => GenerationModel::Flux2LoraAngles, - TauriImageModel::FluxDevJuggernaut => GenerationModel::FluxDevJuggernaut, - TauriImageModel::FluxPro1 => GenerationModel::FluxPro1, - TauriImageModel::FluxPro11 => GenerationModel::FluxPro11, - TauriImageModel::FluxPro11Ultra => GenerationModel::FluxPro11Ultra, - TauriImageModel::FluxProKontextMax => GenerationModel::FluxProKontextMax, - TauriImageModel::Gemini25Flash => GenerationModel::NanoBanana, // NB: Alias for Nano Banana - TauriImageModel::GptImage1 => GenerationModel::GptImage1, - TauriImageModel::GptImage1p5 => GenerationModel::GptImage1p5, - TauriImageModel::GptImage2 => GenerationModel::GptImage2, - TauriImageModel::GrokImage => GenerationModel::GrokImage, - TauriImageModel::NanoBanana => GenerationModel::NanoBanana, - TauriImageModel::NanoBanana2 => GenerationModel::NanoBanana2, - TauriImageModel::NanoBananaPro => GenerationModel::NanoBananaPro, - TauriImageModel::QwenEdit2511Angles => GenerationModel::QwenEdit2511Angles, - TauriImageModel::Seedream4 => GenerationModel::Seedream4, - TauriImageModel::Seedream4p5 => GenerationModel::Seedream4p5, - TauriImageModel::Seedream5Lite => GenerationModel::Seedream5Lite, - // TODO: - TauriImageModel::Midjourney => GenerationModel::Flux1Dev, // Fallback - TauriImageModel::Recraft3 => GenerationModel::Flux1Dev, // Fallback - } -} diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/mod.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/mod.rs index af712f0cb7..657bbe3d3f 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/mod.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/mod.rs @@ -1 +1,6 @@ -pub mod handle_artcraft_via_omni_endpoint; +mod handle_artcraft; +mod handle_artcraft_via_legacy; +mod handle_artcraft_via_omni_endpoint; +mod model_mapping; + +pub use handle_artcraft::handle_artcraft; 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 new file mode 100644 index 0000000000..fa6fee7e07 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/model_mapping.rs @@ -0,0 +1,79 @@ +use artcraft_router::api::common_image_model::CommonImageModel as RouterCommonImageModel; +use enums::common::generation::common_image_model::CommonImageModel as OmniCommonImageModel; + +use crate::core::commands::generate::generate_image::tauri_image_model::TauriImageModel; +use crate::core::events::generation_events::common::GenerationModel; + +/// Map TauriImageModel to the omni endpoint's CommonImageModel. +/// Returns None for models not supported by the omni endpoint. +pub fn map_to_omni_image_model(model: TauriImageModel) -> Option { + match model { + TauriImageModel::Flux1Dev => Some(OmniCommonImageModel::Flux1Dev), + TauriImageModel::Flux1Schnell => Some(OmniCommonImageModel::Flux1Schnell), + TauriImageModel::FluxPro11 => Some(OmniCommonImageModel::FluxPro11), + TauriImageModel::FluxPro11Ultra => Some(OmniCommonImageModel::FluxPro11Ultra), + TauriImageModel::GptImage1 => Some(OmniCommonImageModel::GptImage1), + TauriImageModel::GptImage1p5 => Some(OmniCommonImageModel::GptImage1p5), + TauriImageModel::GptImage2 => Some(OmniCommonImageModel::GptImage2), + TauriImageModel::NanoBanana => Some(OmniCommonImageModel::NanoBanana), + TauriImageModel::NanoBanana2 => Some(OmniCommonImageModel::NanoBanana2), + TauriImageModel::NanoBananaPro => Some(OmniCommonImageModel::NanoBananaPro), + TauriImageModel::Gemini25Flash => Some(OmniCommonImageModel::NanoBanana), + TauriImageModel::Seedream4 => Some(OmniCommonImageModel::Seedream4), + TauriImageModel::Seedream4p5 => Some(OmniCommonImageModel::Seedream4p5), + TauriImageModel::Seedream5Lite => Some(OmniCommonImageModel::Seedream5Lite), + _ => None, + } +} + +/// Map TauriImageModel to the artcraft_router's CommonImageModel. +/// Returns None for models not supported by the router (Grok, Midjourney). +pub fn map_to_router_image_model(model: TauriImageModel) -> 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::QwenEdit2511Angles => Some(RouterCommonImageModel::QwenEdit2511Angles), + TauriImageModel::Flux2LoraAngles => Some(RouterCommonImageModel::Flux2LoraAngles), + _ => None, + } +} + +/// Map TauriImageModel to GenerationModel for frontend events. +pub fn map_to_generation_model(model: TauriImageModel) -> GenerationModel { + match model { + TauriImageModel::Flux1Dev => GenerationModel::Flux1Dev, + TauriImageModel::Flux1Schnell => GenerationModel::Flux1Schnell, + TauriImageModel::FluxPro11 => GenerationModel::FluxPro11, + TauriImageModel::FluxPro11Ultra => GenerationModel::FluxPro11Ultra, + TauriImageModel::GrokImage => GenerationModel::GrokImage, + TauriImageModel::Recraft3 => GenerationModel::Flux1Dev, // Fallback + TauriImageModel::GptImage1 => GenerationModel::GptImage1, + TauriImageModel::GptImage1p5 => GenerationModel::GptImage1p5, + TauriImageModel::GptImage2 => GenerationModel::GptImage2, + TauriImageModel::Gemini25Flash => GenerationModel::NanoBanana, + TauriImageModel::NanoBanana => GenerationModel::NanoBanana, + TauriImageModel::NanoBanana2 => GenerationModel::NanoBanana2, + TauriImageModel::NanoBananaPro => GenerationModel::NanoBananaPro, + TauriImageModel::Seedream4 => GenerationModel::Seedream4, + TauriImageModel::Seedream4p5 => GenerationModel::Seedream4p5, + TauriImageModel::Seedream5Lite => GenerationModel::Seedream5Lite, + TauriImageModel::Midjourney => GenerationModel::Flux1Dev, // Fallback + TauriImageModel::FluxProKontextMax => GenerationModel::FluxProKontextMax, + TauriImageModel::QwenEdit2511Angles => GenerationModel::QwenEdit2511Angles, + TauriImageModel::Flux2LoraAngles => GenerationModel::Flux2LoraAngles, + TauriImageModel::FluxDevJuggernaut => GenerationModel::FluxDevJuggernaut, + TauriImageModel::FluxPro1 => GenerationModel::FluxPro1, + } +}