From ea47bab95070305b2f395652aa8be12c22d7dbcd Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 4 Mar 2026 04:31:48 -0500 Subject: [PATCH] tests updated --- .../src/api/common_image_model.rs | 8 ++ .../artcraft_router/src/client/mod.rs | 1 + ...estimate_image_cost_artcraft_flux_1_dev.rs | 47 +++++++--- .../artcraft_router/src/test_helpers.rs | 66 ++++++++++++++ .../handle_image_edit_artcraft_via_router.rs | 1 + ...andle_artcraft_flux_1_dev_text_to_image.rs | 91 ------------------- ...e_artcraft_flux_1_schnell_text_to_image.rs | 90 ------------------ ...dle_artcraft_flux_pro_1p1_text_to_image.rs | 90 ------------------ ...tcraft_flux_pro_1p1_ultra_text_to_image.rs | 91 ------------------- .../artcraft/handle_text_to_image_artcraft.rs | 14 +-- ...andle_text_to_image_artcraft_via_router.rs | 1 + .../enqueue/text_to_image/artcraft/mod.rs | 6 +- .../image/estimate_image_cost_handler.rs | 1 + 13 files changed, 118 insertions(+), 389 deletions(-) delete mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_1_dev_text_to_image.rs delete mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_1_schnell_text_to_image.rs delete mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_pro_1p1_text_to_image.rs delete mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_pro_1p1_ultra_text_to_image.rs diff --git a/crates/api_clients/artcraft/artcraft_router/src/api/common_image_model.rs b/crates/api_clients/artcraft/artcraft_router/src/api/common_image_model.rs index 64fa26beca..9a8bbfeabd 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/api/common_image_model.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/api/common_image_model.rs @@ -5,6 +5,14 @@ use serde_derive::{Deserialize, Serialize}; #[derive(Copy, Clone, Debug, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum CommonImageModel { + #[serde(rename = "flux_1_dev")] + Flux1Dev, + #[serde(rename = "flux_1_schnell")] + Flux1Schnell, + #[serde(rename = "flux_pro_1p1")] + FluxPro11, + #[serde(rename = "flux_pro_1p1_ultra")] + FluxPro11Ultra, #[serde(rename = "gpt_image_1p5")] GptImage1p5, #[serde(rename = "nano_banana")] diff --git a/crates/api_clients/artcraft/artcraft_router/src/client/mod.rs b/crates/api_clients/artcraft/artcraft_router/src/client/mod.rs index aebd19c5af..8a30739830 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/client/mod.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/client/mod.rs @@ -1,3 +1,4 @@ +pub mod generation_mode_mismatch_strategy; pub mod multi_router_client; pub mod multi_router_client_builder; pub mod request_mismatch_mitigation_strategy; diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/cost/artcraft/estimate_image_cost_artcraft_flux_1_dev.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/cost/artcraft/estimate_image_cost_artcraft_flux_1_dev.rs index fffd502cc9..e5f0376b32 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/cost/artcraft/estimate_image_cost_artcraft_flux_1_dev.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/cost/artcraft/estimate_image_cost_artcraft_flux_1_dev.rs @@ -1,15 +1,28 @@ +use artcraft_api_defs::generate::image::text::generate_flux_1_dev_text_to_image::GenerateFlux1DevTextToImageNumImages; + use crate::generate::generate_image::image_generation_cost_estimate::ImageGenerationCostEstimate; use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_flux_1_dev::PlanArtcraftFlux1Dev; pub(crate) fn estimate_image_cost_artcraft_flux_1_dev( - _plan: &PlanArtcraftFlux1Dev<'_>, + plan: &PlanArtcraftFlux1Dev<'_>, ) -> ImageGenerationCostEstimate { - // Pricing: not available + // Pricing: $0.02/image (2 cents). 1 credit = 1 USD cent. + let cost_per_image: u64 = 2; + + let num_images: u64 = match plan.num_images { + GenerateFlux1DevTextToImageNumImages::One => 1, + GenerateFlux1DevTextToImageNumImages::Two => 2, + GenerateFlux1DevTextToImageNumImages::Three => 3, + GenerateFlux1DevTextToImageNumImages::Four => 4, + }; + + let cost_in_usd_cents = cost_per_image * num_images; + ImageGenerationCostEstimate { - cost_in_credits: None, - cost_in_usd_cents: None, - is_free: true, - is_unlimited: true, + cost_in_credits: Some(cost_in_usd_cents), + cost_in_usd_cents: Some(cost_in_usd_cents), + is_free: false, + is_unlimited: false, is_rate_limited: false, has_watermark: false, } @@ -22,8 +35,7 @@ mod tests { use crate::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy; use crate::generate::generate_image::generate_image_request::GenerateImageRequest; - #[test] - fn test_estimate_cost_is_none() { + fn estimate_usd_cents(image_batch_count: u16) -> u64 { let request = GenerateImageRequest { model: CommonImageModel::Flux1Dev, provider: Provider::Artcraft, @@ -31,15 +43,24 @@ mod tests { image_inputs: None, resolution: None, aspect_ratio: None, - image_batch_count: Some(1), + image_batch_count: Some(image_batch_count), request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::ErrorOut, generation_mode_mismatch_strategy: None, idempotency_token: None, }; - let estimate = request.build() + request.build() .expect("build should succeed") - .estimate_costs(); - assert!(estimate.cost_in_usd_cents.is_none()); - assert!(estimate.cost_in_credits.is_none()); + .estimate_costs() + .cost_in_usd_cents + .expect("cost_in_usd_cents should be present") + } + + #[test] + fn test_estimate_cost_usd_cents() { + // $0.02/image = 2 cents each + assert_eq!(estimate_usd_cents(1), 2); + assert_eq!(estimate_usd_cents(2), 4); + assert_eq!(estimate_usd_cents(3), 6); + assert_eq!(estimate_usd_cents(4), 8); } } diff --git a/crates/api_clients/artcraft/artcraft_router/src/test_helpers.rs b/crates/api_clients/artcraft/artcraft_router/src/test_helpers.rs index 9a5b2c1972..a2ba52c60d 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/test_helpers.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/test_helpers.rs @@ -31,6 +31,7 @@ pub fn base_image_request() -> GenerateImageRequest<'static> { aspect_ratio: None, image_batch_count: None, request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::ErrorOut, + generation_mode_mismatch_strategy: None, idempotency_token: None, } } @@ -53,6 +54,7 @@ pub fn base_fal_image_request() -> GenerateImageRequest<'static> { aspect_ratio: None, image_batch_count: None, request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::ErrorOut, + generation_mode_mismatch_strategy: None, idempotency_token: None, } } @@ -67,6 +69,7 @@ pub fn base_seedream_4_image_request() -> GenerateImageRequest<'static> { aspect_ratio: None, image_batch_count: None, request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::ErrorOut, + generation_mode_mismatch_strategy: None, idempotency_token: None, } } @@ -81,6 +84,7 @@ pub fn base_seedream_4p5_image_request() -> GenerateImageRequest<'static> { aspect_ratio: None, image_batch_count: None, request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::ErrorOut, + generation_mode_mismatch_strategy: None, idempotency_token: None, } } @@ -95,6 +99,7 @@ pub fn base_nano_banana_image_request() -> GenerateImageRequest<'static> { aspect_ratio: None, image_batch_count: None, request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::ErrorOut, + generation_mode_mismatch_strategy: None, idempotency_token: None, } } @@ -109,6 +114,67 @@ pub fn base_gpt_image_1p5_image_request() -> GenerateImageRequest<'static> { aspect_ratio: None, image_batch_count: None, request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::ErrorOut, + generation_mode_mismatch_strategy: None, + idempotency_token: None, + } +} + +pub fn base_flux_1_dev_image_request() -> GenerateImageRequest<'static> { + GenerateImageRequest { + model: CommonImageModel::Flux1Dev, + provider: Provider::Artcraft, + prompt: Some("a cat in space"), + image_inputs: None, + resolution: None, + aspect_ratio: None, + image_batch_count: None, + request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::ErrorOut, + generation_mode_mismatch_strategy: None, + idempotency_token: None, + } +} + +pub fn base_flux_1_schnell_image_request() -> GenerateImageRequest<'static> { + GenerateImageRequest { + model: CommonImageModel::Flux1Schnell, + provider: Provider::Artcraft, + prompt: Some("a cat in space"), + image_inputs: None, + resolution: None, + aspect_ratio: None, + image_batch_count: None, + request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::ErrorOut, + generation_mode_mismatch_strategy: None, + idempotency_token: None, + } +} + +pub fn base_flux_pro_1p1_image_request() -> GenerateImageRequest<'static> { + GenerateImageRequest { + model: CommonImageModel::FluxPro11, + provider: Provider::Artcraft, + prompt: Some("a cat in space"), + image_inputs: None, + resolution: None, + aspect_ratio: None, + image_batch_count: None, + request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::ErrorOut, + generation_mode_mismatch_strategy: None, + idempotency_token: None, + } +} + +pub fn base_flux_pro_1p1_ultra_image_request() -> GenerateImageRequest<'static> { + GenerateImageRequest { + model: CommonImageModel::FluxPro11Ultra, + provider: Provider::Artcraft, + prompt: Some("a cat in space"), + image_inputs: None, + resolution: None, + aspect_ratio: None, + image_batch_count: None, + request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::ErrorOut, + generation_mode_mismatch_strategy: None, idempotency_token: None, } } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/artcraft/handle_image_edit_artcraft_via_router.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/artcraft/handle_image_edit_artcraft_via_router.rs index b1e7fc3b1f..505abc226b 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/artcraft/handle_image_edit_artcraft_via_router.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/artcraft/handle_image_edit_artcraft_via_router.rs @@ -74,6 +74,7 @@ pub(super) async fn handle_image_edit_artcraft_via_router( aspect_ratio, image_batch_count: request.image_count.map(|n| n as u16), request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade, + generation_mode_mismatch_strategy: None, idempotency_token: None, }; diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_1_dev_text_to_image.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_1_dev_text_to_image.rs deleted file mode 100644 index 85410dc6d3..0000000000 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_1_dev_text_to_image.rs +++ /dev/null @@ -1,91 +0,0 @@ -use crate::core::api_adapters::aspect_ratio::convert::aspect_ratio_to_artcraft_flux_1_dev::aspect_ratio_to_artcraft_flux_1_dev; -use crate::core::api_adapters::aspect_ratio::convert::aspect_ratio_to_artcraft_nano_banana_pro::aspect_ratio_to_artcraft_nano_banana_pro; -use crate::core::commands::enqueue::generate_error::GenerateError; -use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; -use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize}; -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; -use artcraft_api_defs::generate::image::multi_function::nano_banana_pro_multi_function_image_gen::NanoBananaProMultiFunctionImageGenAspectRatio; -use artcraft_api_defs::generate::image::text::generate_flux_1_dev_text_to_image::{GenerateFlux1DevTextToImageAspectRatio, GenerateFlux1DevTextToImageNumImages, GenerateFlux1DevTextToImageRequest}; -use enums::common::generation_provider::GenerationProvider; -use enums::tauri::tasks::task_type::TaskType; -use idempotency::uuid::generate_random_uuid; -use log::{error, info}; -use artcraft_client::endpoints::generate::image::text::generate_flux_1_dev_text_to_image::generate_flux_1_dev_text_to_image; - -pub async fn handle_artcraft_flux_1_dev_text_to_image( - request: &EnqueueTextToImageRequest, - app_env_configs: &AppEnvConfigs, - storyteller_creds_manager: &StorytellerCredentialManager, -) -> Result { - let creds = match storyteller_creds_manager.get_credentials()? { - Some(creds) => creds, - None => { - return Err(GenerateError::needs_storyteller_credentials()); - }, - }; - - info!("enqueue Flux 1 Dev"); - - let uuid_idempotency_token = generate_random_uuid(); - - let request = GenerateFlux1DevTextToImageRequest { - uuid_idempotency_token, - prompt: request.prompt.clone(), - aspect_ratio: get_aspect_ratio(request), - num_images: request.number_images - .and_then(|num| match num { - 1 => Some(GenerateFlux1DevTextToImageNumImages::One), - 2 => Some(GenerateFlux1DevTextToImageNumImages::Two), - 3 => Some(GenerateFlux1DevTextToImageNumImages::Three), - 4 => Some(GenerateFlux1DevTextToImageNumImages::Four), - _ => None, // NB: use service defaults - }), - }; - - let result = generate_flux_1_dev_text_to_image( - &app_env_configs.storyteller_host, - Some(&creds), - request, - ).await; - - let job_token = match result { - Ok(enqueued) => { - info!("Successfully enqueued Artcraft Flux 1 Dev text to image generation"); - enqueued.inference_job_token - } - Err(err) => { - error!("Failed to use Artcraft Flux 1 Dev text to image generation: {:?}", err); - return Err(GenerateError::from(err)); - } - }; - - Ok(TaskEnqueueSuccess { - provider: GenerationProvider::Artcraft, - task_type: TaskType::ImageGeneration, - model: Some(GenerationModel::Flux1Dev), - provider_job_id: Some(job_token.to_string()), - }) -} - -fn get_aspect_ratio(request: &EnqueueTextToImageRequest) -> Option { - if let Some(common_aspect_ratio) = request.common_aspect_ratio { - // Handle modern aspect ratio - let aspect = aspect_ratio_to_artcraft_flux_1_dev(common_aspect_ratio); - return Some(aspect); - } - - if let Some(aspect_ratio) = request.aspect_ratio { - // Handle deprecated aspect ratio - let aspect = match aspect_ratio { - TextToImageSize::Tall => GenerateFlux1DevTextToImageAspectRatio::PortraitNineBySixteen, - TextToImageSize::Wide => GenerateFlux1DevTextToImageAspectRatio::LandscapeSixteenByNine, - TextToImageSize::Auto => GenerateFlux1DevTextToImageAspectRatio::SquareHd, - TextToImageSize::Square => GenerateFlux1DevTextToImageAspectRatio::SquareHd, - }; - return Some(aspect); - } - - None -} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_1_schnell_text_to_image.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_1_schnell_text_to_image.rs deleted file mode 100644 index f9107cda0a..0000000000 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_1_schnell_text_to_image.rs +++ /dev/null @@ -1,90 +0,0 @@ -use crate::core::api_adapters::aspect_ratio::convert::aspect_ratio_to_artcraft_flux_1_dev::aspect_ratio_to_artcraft_flux_1_dev; -use crate::core::api_adapters::aspect_ratio::convert::aspect_ratio_to_artcraft_flux_1_schnell::aspect_ratio_to_artcraft_flux_1_schnell; -use crate::core::commands::enqueue::generate_error::GenerateError; -use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; -use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize}; -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; -use artcraft_api_defs::generate::image::text::generate_flux_1_schnell_text_to_image::{GenerateFlux1SchnellTextToImageAspectRatio, GenerateFlux1SchnellTextToImageNumImages, GenerateFlux1SchnellTextToImageRequest}; -use enums::common::generation_provider::GenerationProvider; -use enums::tauri::tasks::task_type::TaskType; -use idempotency::uuid::generate_random_uuid; -use log::{error, info}; -use artcraft_client::endpoints::generate::image::text::generate_flux_1_schnell_text_to_image::generate_flux_1_schnell_text_to_image; - -pub async fn handle_artcraft_flux_1_schnell_text_to_image( - request: &EnqueueTextToImageRequest, - app_env_configs: &AppEnvConfigs, - storyteller_creds_manager: &StorytellerCredentialManager, -) -> Result { - let creds = match storyteller_creds_manager.get_credentials()? { - Some(creds) => creds, - None => { - return Err(GenerateError::needs_storyteller_credentials()); - }, - }; - - info!("enqueue Flux 1 Schnell"); - - let uuid_idempotency_token = generate_random_uuid(); - - let request = GenerateFlux1SchnellTextToImageRequest { - uuid_idempotency_token, - prompt: request.prompt.clone(), - aspect_ratio: get_aspect_ratio(request), - num_images: request.number_images - .and_then(|num| match num { - 1 => Some(GenerateFlux1SchnellTextToImageNumImages::One), - 2 => Some(GenerateFlux1SchnellTextToImageNumImages::Two), - 3 => Some(GenerateFlux1SchnellTextToImageNumImages::Three), - 4 => Some(GenerateFlux1SchnellTextToImageNumImages::Four), - _ => None, // NB: use service defaults - }), - }; - - let result = generate_flux_1_schnell_text_to_image( - &app_env_configs.storyteller_host, - Some(&creds), - request, - ).await; - - let job_token = match result { - Ok(enqueued) => { - info!("Successfully enqueued Artcraft Flux 1 Schnell text to image generation"); - enqueued.inference_job_token - } - Err(err) => { - error!("Failed to use Artcraft Flux 1 Schnell text to image generation: {:?}", err); - return Err(GenerateError::from(err)); - } - }; - - Ok(TaskEnqueueSuccess { - provider: GenerationProvider::Artcraft, - task_type: TaskType::ImageGeneration, - model: Some(GenerationModel::Flux1Schnell), - provider_job_id: Some(job_token.to_string()), - }) -} - -fn get_aspect_ratio(request: &EnqueueTextToImageRequest) -> Option { - if let Some(common_aspect_ratio) = request.common_aspect_ratio { - // Handle modern aspect ratio - let aspect = aspect_ratio_to_artcraft_flux_1_schnell(common_aspect_ratio); - return Some(aspect); - } - - if let Some(aspect_ratio) = request.aspect_ratio { - // Handle deprecated aspect ratio - let aspect = match aspect_ratio { - TextToImageSize::Tall => GenerateFlux1SchnellTextToImageAspectRatio::PortraitNineBySixteen, - TextToImageSize::Wide => GenerateFlux1SchnellTextToImageAspectRatio::LandscapeSixteenByNine, - TextToImageSize::Auto => GenerateFlux1SchnellTextToImageAspectRatio::SquareHd, - TextToImageSize::Square => GenerateFlux1SchnellTextToImageAspectRatio::SquareHd, - }; - return Some(aspect); - } - - None -} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_pro_1p1_text_to_image.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_pro_1p1_text_to_image.rs deleted file mode 100644 index ef62d8931c..0000000000 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_pro_1p1_text_to_image.rs +++ /dev/null @@ -1,90 +0,0 @@ -use crate::core::api_adapters::aspect_ratio::convert::aspect_ratio_to_artcraft_flux_1_schnell::aspect_ratio_to_artcraft_flux_1_schnell; -use crate::core::api_adapters::aspect_ratio::convert::aspect_ratio_to_artcraft_flux_pro_1p1::aspect_ratio_to_artcraft_flux_pro_1p1; -use crate::core::commands::enqueue::generate_error::GenerateError; -use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; -use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize}; -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; -use artcraft_api_defs::generate::image::text::generate_flux_pro_11_text_to_image::{GenerateFluxPro11TextToImageAspectRatio, GenerateFluxPro11TextToImageNumImages, GenerateFluxPro11TextToImageRequest}; -use enums::common::generation_provider::GenerationProvider; -use enums::tauri::tasks::task_type::TaskType; -use idempotency::uuid::generate_random_uuid; -use log::{error, info}; -use artcraft_client::endpoints::generate::image::text::generate_flux_pro_11_text_to_image::generate_flux_pro_11_text_to_image; - -pub async fn handle_artcraft_flux_pro_1p1_text_to_image( - request: &EnqueueTextToImageRequest, - app_env_configs: &AppEnvConfigs, - storyteller_creds_manager: &StorytellerCredentialManager, -) -> Result { - let creds = match storyteller_creds_manager.get_credentials()? { - Some(creds) => creds, - None => { - return Err(GenerateError::needs_storyteller_credentials()); - }, - }; - - info!("enqueue Flux Pro 1.1"); - - let uuid_idempotency_token = generate_random_uuid(); - - let request = GenerateFluxPro11TextToImageRequest { - uuid_idempotency_token, - prompt: request.prompt.clone(), - aspect_ratio: get_aspect_ratio(request), - num_images: request.number_images - .and_then(|num| match num { - 1 => Some(GenerateFluxPro11TextToImageNumImages::One), - 2 => Some(GenerateFluxPro11TextToImageNumImages::Two), - 3 => Some(GenerateFluxPro11TextToImageNumImages::Three), - 4 => Some(GenerateFluxPro11TextToImageNumImages::Four), - _ => None, // NB: use service defaults - }), - }; - - let result = generate_flux_pro_11_text_to_image( - &app_env_configs.storyteller_host, - Some(&creds), - request, - ).await; - - let job_token = match result { - Ok(enqueued) => { - info!("Successfully enqueued Artcraft Flux Pro 1.1 text to image generation"); - enqueued.inference_job_token - } - Err(err) => { - error!("Failed to use Artcraft Flux Pro 1.1 text to image generation: {:?}", err); - return Err(GenerateError::from(err)); - } - }; - - Ok(TaskEnqueueSuccess { - provider: GenerationProvider::Artcraft, - task_type: TaskType::ImageGeneration, - model: Some(GenerationModel::FluxPro11), - provider_job_id: Some(job_token.to_string()), - }) -} - -fn get_aspect_ratio(request: &EnqueueTextToImageRequest) -> Option { - if let Some(common_aspect_ratio) = request.common_aspect_ratio { - // Handle modern aspect ratio - let aspect = aspect_ratio_to_artcraft_flux_pro_1p1(common_aspect_ratio); - return Some(aspect); - } - - if let Some(aspect_ratio) = request.aspect_ratio { - // Handle deprecated aspect ratio - let aspect = match aspect_ratio { - TextToImageSize::Tall => GenerateFluxPro11TextToImageAspectRatio::PortraitNineBySixteen, - TextToImageSize::Wide => GenerateFluxPro11TextToImageAspectRatio::LandscapeSixteenByNine, - TextToImageSize::Square => GenerateFluxPro11TextToImageAspectRatio::Square, - TextToImageSize::Auto => GenerateFluxPro11TextToImageAspectRatio::Square, - }; - return Some(aspect); - } - - None -} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_pro_1p1_ultra_text_to_image.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_pro_1p1_ultra_text_to_image.rs deleted file mode 100644 index 5876ada04c..0000000000 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_artcraft_flux_pro_1p1_ultra_text_to_image.rs +++ /dev/null @@ -1,91 +0,0 @@ -use crate::core::api_adapters::aspect_ratio::convert::aspect_ratio_to_artcraft_flux_pro_1p1::aspect_ratio_to_artcraft_flux_pro_1p1; -use crate::core::api_adapters::aspect_ratio::convert::aspect_ratio_to_artcraft_flux_pro_1p1_ultra::aspect_ratio_to_artcraft_flux_pro_1p1_ultra; -use crate::core::commands::enqueue::generate_error::GenerateError; -use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; -use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize}; -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; -use artcraft_api_defs::generate::image::text::generate_flux_pro_11_ultra_text_to_image::{GenerateFluxPro11UltraTextToImageAspectRatio, GenerateFluxPro11UltraTextToImageNumImages, GenerateFluxPro11UltraTextToImageRequest}; -use enums::common::generation_provider::GenerationProvider; -use enums::tauri::tasks::task_type::TaskType; -use idempotency::uuid::generate_random_uuid; -use log::{error, info}; -use artcraft_client::endpoints::generate::image::text::generate_flux_pro_11_ultra_text_to_image::generate_flux_pro_11_ultra_text_to_image; - -pub async fn handle_artcraft_flux_pro_1p1_ultra_text_to_image( - request: &EnqueueTextToImageRequest, - app_env_configs: &AppEnvConfigs, - storyteller_creds_manager: &StorytellerCredentialManager, -) -> Result { - - let creds = match storyteller_creds_manager.get_credentials()? { - Some(creds) => creds, - None => { - return Err(GenerateError::needs_storyteller_credentials()); - }, - }; - - info!("enqueue Flux Pro 1.1 Ultra"); - - let uuid_idempotency_token = generate_random_uuid(); - - let request = GenerateFluxPro11UltraTextToImageRequest { - uuid_idempotency_token, - prompt: request.prompt.clone(), - aspect_ratio: get_aspect_ratio(request), - num_images: request.number_images - .and_then(|num| match num { - 1 => Some(GenerateFluxPro11UltraTextToImageNumImages::One), - 2 => Some(GenerateFluxPro11UltraTextToImageNumImages::Two), - 3 => Some(GenerateFluxPro11UltraTextToImageNumImages::Three), - 4 => Some(GenerateFluxPro11UltraTextToImageNumImages::Four), - _ => None, // NB: use service defaults - }), - }; - - let result = generate_flux_pro_11_ultra_text_to_image( - &app_env_configs.storyteller_host, - Some(&creds), - request, - ).await; - - let job_token = match result { - Ok(enqueued) => { - info!("Successfully enqueued Artcraft Flux Pro 1.1 Ultra text to image generation"); - enqueued.inference_job_token - } - Err(err) => { - error!("Failed to use Artcraft Flux Pro 1.1 Ultra text to image generation: {:?}", err); - return Err(GenerateError::from(err)); - } - }; - - Ok(TaskEnqueueSuccess { - provider: GenerationProvider::Artcraft, - task_type: TaskType::ImageGeneration, - model: Some(GenerationModel::FluxPro11Ultra), - provider_job_id: Some(job_token.to_string()), - }) -} - -fn get_aspect_ratio(request: &EnqueueTextToImageRequest) -> Option { - if let Some(common_aspect_ratio) = request.common_aspect_ratio { - // Handle modern aspect ratio - let aspect = aspect_ratio_to_artcraft_flux_pro_1p1_ultra(common_aspect_ratio); - return Some(aspect); - } - - if let Some(aspect_ratio) = request.aspect_ratio { - // Handle deprecated aspect ratio - let aspect = match aspect_ratio { - TextToImageSize::Tall => GenerateFluxPro11UltraTextToImageAspectRatio::PortraitNineBySixteen, - TextToImageSize::Wide => GenerateFluxPro11UltraTextToImageAspectRatio::LandscapeSixteenByNine, - TextToImageSize::Square => GenerateFluxPro11UltraTextToImageAspectRatio::Square, - TextToImageSize::Auto => GenerateFluxPro11UltraTextToImageAspectRatio::Square, - }; - return Some(aspect); - } - - None -} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_text_to_image_artcraft.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_text_to_image_artcraft.rs index f593d89739..5a382e58fb 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_text_to_image_artcraft.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_text_to_image_artcraft.rs @@ -1,9 +1,5 @@ use crate::core::commands::enqueue::generate_error::GenerateError; use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; -use crate::core::commands::enqueue::text_to_image::artcraft::handle_artcraft_flux_1_dev_text_to_image::handle_artcraft_flux_1_dev_text_to_image; -use crate::core::commands::enqueue::text_to_image::artcraft::handle_artcraft_flux_1_schnell_text_to_image::handle_artcraft_flux_1_schnell_text_to_image; -use crate::core::commands::enqueue::text_to_image::artcraft::handle_artcraft_flux_pro_1p1_text_to_image::handle_artcraft_flux_pro_1p1_text_to_image; -use crate::core::commands::enqueue::text_to_image::artcraft::handle_artcraft_flux_pro_1p1_ultra_text_to_image::handle_artcraft_flux_pro_1p1_ultra_text_to_image; use crate::core::commands::enqueue::text_to_image::artcraft::handle_artcraft_gpt_image_1_text_to_image::handle_artcraft_gpt_image_1_text_to_image; use crate::core::commands::enqueue::text_to_image::artcraft::handle_text_to_image_artcraft_via_router::handle_text_to_image_artcraft_via_router; use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageModel}; @@ -24,7 +20,7 @@ pub async fn handle_text_to_image_artcraft( app_env_configs: &AppEnvConfigs, storyteller_creds_manager: &StorytellerCredentialManager, ) -> Result { - + match model { TextToImageModel::Midjourney | TextToImageModel::GrokImage => { Err(GenerateError::BadProviderForModel { @@ -35,10 +31,10 @@ pub async fn handle_text_to_image_artcraft( TextToImageModel::Recraft3 => { Err(GenerateError::NotYetImplemented(format!("not yet implemented in Artcraft"))) } - TextToImageModel::Flux1Dev => handle_artcraft_flux_1_dev_text_to_image(request, app_env_configs, storyteller_creds_manager).await, - TextToImageModel::Flux1Schnell => handle_artcraft_flux_1_schnell_text_to_image(request, app_env_configs, storyteller_creds_manager).await, - TextToImageModel::FluxPro11 =>handle_artcraft_flux_pro_1p1_text_to_image(request, app_env_configs, storyteller_creds_manager).await, - TextToImageModel::FluxPro11Ultra => handle_artcraft_flux_pro_1p1_ultra_text_to_image(request, app_env_configs, storyteller_creds_manager).await, + TextToImageModel::Flux1Dev => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::Flux1Dev, GenerationModel::Flux1Dev).await, + TextToImageModel::Flux1Schnell => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::Flux1Schnell, GenerationModel::Flux1Schnell).await, + TextToImageModel::FluxPro11 => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::FluxPro11, GenerationModel::FluxPro11).await, + TextToImageModel::FluxPro11Ultra => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::FluxPro11Ultra, GenerationModel::FluxPro11Ultra).await, TextToImageModel::GptImage1 => handle_artcraft_gpt_image_1_text_to_image(request, app_env_configs, storyteller_creds_manager).await, TextToImageModel::GptImage1p5 => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::GptImage1p5, GenerationModel::GptImage1p5).await, TextToImageModel::Gemini25Flash | TextToImageModel::NanoBanana => { diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_text_to_image_artcraft_via_router.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_text_to_image_artcraft_via_router.rs index bc27b82242..6063ff48e8 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_text_to_image_artcraft_via_router.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/handle_text_to_image_artcraft_via_router.rs @@ -52,6 +52,7 @@ pub(super) async fn handle_text_to_image_artcraft_via_router( aspect_ratio, image_batch_count: request.number_images.map(|n| n as u16), request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade, + generation_mode_mismatch_strategy: None, idempotency_token: None, }; diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/mod.rs index 2671f4be65..572b124965 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/mod.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/artcraft/mod.rs @@ -1,8 +1,4 @@ pub (super) mod handle_text_to_image_artcraft; mod handle_text_to_image_artcraft_via_router; -mod handle_artcraft_flux_1_dev_text_to_image; -mod handle_artcraft_flux_1_schnell_text_to_image; -mod handle_artcraft_flux_pro_1p1_text_to_image; -mod handle_artcraft_flux_pro_1p1_ultra_text_to_image; -mod handle_artcraft_gpt_image_1_text_to_image; \ No newline at end of file +mod handle_artcraft_gpt_image_1_text_to_image; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs index 8e9f483a41..f54ce474ff 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs @@ -46,6 +46,7 @@ pub async fn estimate_image_cost_handler( aspect_ratio: router_aspect_ratio, image_batch_count: request.image_batch_count, request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayLessDowngrade, + generation_mode_mismatch_strategy: None, idempotency_token: None, };