From bb5e4d28232e48b12cf2c5d40617d59ed680de8d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 31 Jul 2025 19:03:22 -0400 Subject: [PATCH] tauri flux kontext --- .../edit/flux_pro_kontext_max_edit_image.rs | 19 +++ .../src/generate/image/edit/mod.rs | 1 + .../enqueue_contextual_edit_image_command.rs | 16 +- .../commands/enqueue/image_edit/errors.rs | 11 ++ .../flux_kontext/handle_flux_kontext_edit.rs | 59 +++++++ .../handle_flux_kontext_edit_artcraft.rs | 148 ++++++++++++++++++ .../enqueue/image_edit/flux_kontext/mod.rs | 2 + .../core/commands/enqueue/image_edit/mod.rs | 3 +- .../enqueue_image_inpaint_command.rs | 18 --- .../handle_flux_pro_1_inpaint_artcraft.rs | 4 +- .../image_inpaint/flux_pro_1_inpaint/mod.rs | 2 +- .../commands/enqueue/task_enqueue_success.rs | 1 + .../core/events/generation_events/common.rs | 2 + .../model/contextual_image_edit_models.rs | 3 + .../enums/src/tauri/tasks/task_model_type.rs | 11 +- 15 files changed, 276 insertions(+), 24 deletions(-) create mode 100644 crates/api_clients/storyteller_client/src/generate/image/edit/flux_pro_kontext_max_edit_image.rs create mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/handle_flux_kontext_edit.rs create mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/handle_flux_kontext_edit_artcraft.rs create mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/mod.rs diff --git a/crates/api_clients/storyteller_client/src/generate/image/edit/flux_pro_kontext_max_edit_image.rs b/crates/api_clients/storyteller_client/src/generate/image/edit/flux_pro_kontext_max_edit_image.rs new file mode 100644 index 0000000000..372179165a --- /dev/null +++ b/crates/api_clients/storyteller_client/src/generate/image/edit/flux_pro_kontext_max_edit_image.rs @@ -0,0 +1,19 @@ +use crate::credentials::storyteller_credential_set::StorytellerCredentialSet; +use crate::error::storyteller_error::StorytellerError; +use crate::utils::api_host::ApiHost; +use crate::utils::basic_json_post_request::basic_json_post_request; +use artcraft_api_defs::generate::image::edit::flux_pro_kontext_max_edit_image::{FluxProKontextMaxEditImageRequest, FluxProKontextMaxEditImageResponse, FLUX_PRO_KONTEXT_MAX_EDIT_IMAGE_PATH}; + + +pub async fn flux_pro_kontext_max_edit_image( + api_host: &ApiHost, + maybe_creds: Option<&StorytellerCredentialSet>, + request: FluxProKontextMaxEditImageRequest, +) -> Result { + Ok(basic_json_post_request( + api_host, + FLUX_PRO_KONTEXT_MAX_EDIT_IMAGE_PATH, + maybe_creds, + request, + ).await?) +} diff --git a/crates/api_clients/storyteller_client/src/generate/image/edit/mod.rs b/crates/api_clients/storyteller_client/src/generate/image/edit/mod.rs index e839328df0..ed56ffa74f 100644 --- a/crates/api_clients/storyteller_client/src/generate/image/edit/mod.rs +++ b/crates/api_clients/storyteller_client/src/generate/image/edit/mod.rs @@ -1 +1,2 @@ +pub mod flux_pro_kontext_max_edit_image; pub mod gpt_image_1_edit_image; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/enqueue_contextual_edit_image_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/enqueue_contextual_edit_image_command.rs index c579f1c7e7..88bf34b801 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/enqueue_contextual_edit_image_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/enqueue_contextual_edit_image_command.rs @@ -43,6 +43,7 @@ use storyteller_client::media_files::get_media_file::get_media_file; use storyteller_client::utils::api_host::ApiHost; use tauri::{AppHandle, Manager, State}; use tokens::tokens::media_files::MediaFileToken; +use crate::core::commands::enqueue::image_edit::flux_kontext::handle_flux_kontext_edit::handle_flux_kontext_edit; #[derive(Deserialize, Debug)] pub struct EnqueueContextualEditImageCommand { @@ -207,6 +208,20 @@ pub async fn handle_request( None => { return Err(InternalContextualEditImageError::NoModelSpecified) } + Some(ContextualImageEditModel::FluxProKontextMax) => { + handle_flux_kontext_edit( + request, + app, + app_data_root, + app_env_configs, + provider_priority_store, + storyteller_creds_manager, + fal_creds_manager, + fal_task_queue, + sora_creds_manager, + sora_task_queue, + ).await? + } Some(ContextualImageEditModel::GptImage1) => { handle_gpt_image_1_edit( request, @@ -221,7 +236,6 @@ pub async fn handle_request( sora_task_queue, ).await? } - // TODO(bt,2025-07-05): Flux Kontext, etc. }; let result = success_event diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/errors.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/errors.rs index 8e46f12817..945522821c 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/errors.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/errors.rs @@ -30,6 +30,12 @@ pub enum InternalContextualEditImageError { requested: u32, }, + // TODO: Maybe just reuse the other error enum variant + /// An invalid number of input images was requested for Flux Kontext. + InvalidNumberOfInputImagesForFluxKontext { + message: String, + }, + /// No sora credentials are set. NeedsSoraCredentials, @@ -84,6 +90,11 @@ impl InternalContextualEditImageError { error_type = EnqueueContextualEditImageErrorType::BadRequest; error_message = format!("Invalid number of input images ({}). Must be between {} and {}", requested, min, max); } + InternalContextualEditImageError::InvalidNumberOfInputImagesForFluxKontext { message } => { + status = CommandErrorStatus::BadRequest; + error_type = EnqueueContextualEditImageErrorType::BadRequest; + error_message = message.clone(); + } InternalContextualEditImageError::SoraError(_) => {} InternalContextualEditImageError::AnyhowError(_) => {} InternalContextualEditImageError::StorytellerError(_) => {} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/handle_flux_kontext_edit.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/handle_flux_kontext_edit.rs new file mode 100644 index 0000000000..cc97c13ce4 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/handle_flux_kontext_edit.rs @@ -0,0 +1,59 @@ +use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::EnqueueContextualEditImageCommand; +use crate::core::commands::enqueue::image_edit::errors::InternalContextualEditImageError; +use crate::core::commands::enqueue::image_edit::flux_kontext::handle_flux_kontext_edit_artcraft::handle_flux_kontext_edit_artcraft; +use crate::core::commands::enqueue::image_edit::gpt_image_1::handle_gpt_image_1_edit_artcraft::handle_gpt_image_1_edit_artcraft; +use crate::core::commands::enqueue::image_edit::gpt_image_1::handle_gpt_image_1_edit_sora::handle_gpt_image_1_edit_sora; +use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::provider_priority::{Provider, ProviderPriorityStore}; +use crate::services::fal::state::fal_credential_manager::FalCredentialManager; +use crate::services::fal::state::fal_task_queue::FalTaskQueue; +use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; +use crate::services::sora::state::sora_task_queue::SoraTaskQueue; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use log::info; +use tauri::AppHandle; + +pub async fn handle_flux_kontext_edit( + request: &EnqueueContextualEditImageCommand, + app: &AppHandle, + app_data_root: &AppDataRoot, + app_env_configs: &AppEnvConfigs, + provider_priority_store: &ProviderPriorityStore, + storyteller_creds_manager: &StorytellerCredentialManager, + fal_creds_manager: &FalCredentialManager, + fal_task_queue: &FalTaskQueue, + sora_creds_manager: &SoraCredentialManager, + sora_task_queue: &SoraTaskQueue, +) -> Result { + + let priority = provider_priority_store.get_priority()?; + + // TODO: Check if providers are available before proceeding. + + info!("Providers by priority: {:?}", priority); + + for provider in priority.iter() { + match provider { + Provider::Fal => { + // Fallthrough (TODO) + } + Provider::Sora => { + // Fallthrough + } + Provider::Artcraft => { + info!("Dispatching gpt-image-1 (edit) via Artcraft..."); + return handle_flux_kontext_edit_artcraft( + request, + app, + app_data_root, + app_env_configs, + storyteller_creds_manager + ).await; + } + } + } + + Err(InternalContextualEditImageError::NoProviderAvailable) +} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/handle_flux_kontext_edit_artcraft.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/handle_flux_kontext_edit_artcraft.rs new file mode 100644 index 0000000000..caa6f4b37f --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/handle_flux_kontext_edit_artcraft.rs @@ -0,0 +1,148 @@ +use crate::core::commands::enqueue::image::enqueue_text_to_image_command::EnqueueTextToImageRequest; +use crate::core::commands::enqueue::image::internal_image_error::InternalImageError; +use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::{EditImageQuality, EditImageSize, EnqueueContextualEditImageCommand}; +use crate::core::commands::enqueue::image_edit::errors::InternalContextualEditImageError; +use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; +use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider}; +use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent; +use crate::core::model::contextual_image_edit_models::ContextualImageEditModel; +use crate::core::model::image_models::ImageModel; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::provider_priority::ProviderPriorityStore; +use crate::services::fal::state::fal_credential_manager::FalCredentialManager; +use crate::services::fal::state::fal_task_queue::FalTaskQueue; +use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; +use crate::services::sora::state::sora_task_queue::SoraTaskQueue; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use anyhow::anyhow; +use artcraft_api_defs::generate::image::edit::flux_pro_kontext_max_edit_image::{FluxProKontextMaxEditImageNumImages, FluxProKontextMaxEditImageRequest}; +use artcraft_api_defs::generate::image::edit::gpt_image_1_edit_image::{GptImage1EditImageImageQuality, GptImage1EditImageImageSize, GptImage1EditImageNumImages, GptImage1EditImageRequest}; +use artcraft_api_defs::generate::image::generate_flux_1_dev_text_to_image::GenerateFlux1DevTextToImageRequest; +use artcraft_api_defs::generate::image::generate_flux_1_schnell_text_to_image::GenerateFlux1SchnellTextToImageRequest; +use artcraft_api_defs::generate::image::generate_flux_pro_11_text_to_image::GenerateFluxPro11TextToImageRequest; +use artcraft_api_defs::generate::image::generate_flux_pro_11_ultra_text_to_image::GenerateFluxPro11UltraTextToImageRequest; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_type::TaskType; +use fal_client::requests::queue::image_gen::enqueue_flux_pro_11_ultra_text_to_image::{enqueue_flux_pro_11_ultra_text_to_image, FluxPro11UltraTextToImageArgs}; +use fal_client::requests::webhook::image::edit::enqueue_flux_pro_kontext_max_edit_webhook::FluxProKontextMaxNumImages; +use fal_client::requests::webhook::image::enqueue_gpt_image_1_edit_image_webhook::GptEditImageQuality; +use idempotency::uuid::generate_random_uuid; +use log::{error, info}; +use storyteller_client::generate::image::edit::flux_pro_kontext_max_edit_image::flux_pro_kontext_max_edit_image; +use storyteller_client::generate::image::edit::gpt_image_1_edit_image::gpt_image_1_edit_image; +use storyteller_client::generate::image::generate_flux_1_dev_text_to_image::generate_flux_1_dev_text_to_image; +use storyteller_client::generate::image::generate_flux_1_schnell_text_to_image::generate_flux_1_schnell_text_to_image; +use storyteller_client::generate::image::generate_flux_pro_11_text_to_image::generate_flux_pro_11_text_to_image; +use storyteller_client::generate::image::generate_flux_pro_11_ultra_text_to_image::generate_flux_pro_11_ultra_text_to_image; +use tauri::AppHandle; + +pub async fn handle_flux_kontext_edit_artcraft( + request: &EnqueueContextualEditImageCommand, + app: &AppHandle, + app_data_root: &AppDataRoot, + app_env_configs: &AppEnvConfigs, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> Result { + + let creds = match storyteller_creds_manager.get_credentials()? { + Some(creds) => creds, + None => { + return Err(InternalContextualEditImageError::NeedsStorytellerCredentials); + }, + }; + + info!("Calling Artcraft flux kontext max edit (edit) ..."); + + let uuid_idempotency_token = generate_random_uuid(); + + let num_images = match request.image_count { + None => None, + Some(1) => Some(FluxProKontextMaxEditImageNumImages::One), + Some(2) => Some(FluxProKontextMaxEditImageNumImages::Two), + Some(3) => Some(FluxProKontextMaxEditImageNumImages::Three), + Some(4) => Some(FluxProKontextMaxEditImageNumImages::Four), + Some(other) => { + return Err(InternalContextualEditImageError::InvalidNumberOfRequestedImages { + min: 1, + max: 4, + requested: other, + }); + }, + }; + + let mut maybe_media_token = None; + + let list_count = request.image_media_tokens + .as_ref() + .map(|tokens| tokens.len()) + .unwrap_or(0); + + if request.scene_image_media_token.is_some() && list_count > 0 { + return Err(InternalContextualEditImageError::InvalidNumberOfInputImagesForFluxKontext { + message: "Cannot specify both scene_image_media_token and image_media_tokens for Flux Kontext Max".to_string() + }); + } + + if list_count > 1 { + return Err(InternalContextualEditImageError::InvalidNumberOfInputImagesForFluxKontext { + message: "Cannot specify more than one image in image_media_tokens for Flux Kontext Max".to_string() + }); + } + + if let Some(media_token) = &request.scene_image_media_token { + maybe_media_token = Some(media_token.clone()); + } + + let maybe_list_media_token = request.image_media_tokens + .as_ref() + .map(|tokens| tokens.first()) + .flatten(); + + if let Some(media_token) = maybe_list_media_token { + maybe_media_token = Some(media_token.clone()); + } + + let media_token = match maybe_media_token { + Some(token) => token, + None => { + return Err(InternalContextualEditImageError::InvalidNumberOfInputImagesForFluxKontext { + message: "No input image specified for Flux Kontext Max edit".to_string(), + }); + }, + }; + + let request = FluxProKontextMaxEditImageRequest { + uuid_idempotency_token, + prompt: Some(request.prompt.clone()), + image_media_token: media_token, + num_images, + }; + + let result = flux_pro_kontext_max_edit_image( + &app_env_configs.storyteller_host, + Some(&creds), + request, + ).await; + + let job_id = match result { + Ok(enqueued) => { + // TODO(bt,2025-07-05): Enqueue job token? + info!("Successfully enqueued Artcraft gpt-image-1. Job token: {}", + enqueued.inference_job_token); + enqueued.inference_job_token + } + Err(err) => { + error!("Failed to use Artcraft gpt-image-1: {:?}", err); + return Err(InternalContextualEditImageError::StorytellerError(err)); + } + }; + + Ok(TaskEnqueueSuccess { + provider: GenerationProvider::Artcraft, + model: Some(GenerationModel::FluxProKontextMax), + provider_job_id: Some(job_id.to_string()), + task_type: TaskType::ImageGeneration, + }) +} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/mod.rs new file mode 100644 index 0000000000..54e79866ca --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/flux_kontext/mod.rs @@ -0,0 +1,2 @@ +pub (super) mod handle_flux_kontext_edit; +mod handle_flux_kontext_edit_artcraft; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/mod.rs index abc8f74a1d..c89e043ec7 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/mod.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/mod.rs @@ -1,3 +1,4 @@ pub mod enqueue_contextual_edit_image_command; pub mod errors; -pub mod gpt_image_1; \ No newline at end of file +pub mod gpt_image_1; +mod flux_kontext; \ No newline at end of file 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 9867d86060..d42ff1c295 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 @@ -80,24 +80,6 @@ pub struct EnqueueInpaintImageCommand { pub image_count: Option, } -#[derive(Deserialize, Debug, Copy, Clone)] -#[serde(rename_all = "snake_case")] -pub enum EditImageSize { - Auto, - Square, - Wide, - Tall, -} - -#[derive(Deserialize, Debug, Copy, Clone)] -#[serde(rename_all = "snake_case")] -pub enum EditImageQuality { - Auto, - High, - Medium, - Low, -} - #[derive(Serialize, Debug)] #[serde(rename_all = "snake_case")] pub enum EnqueueInpaintImageErrorType { 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 8600e93bbc..8522341cbd 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 @@ -103,7 +103,7 @@ pub async fn handle_flux_pro_1_inpaint_artcraft( let job_id = match result { Ok(enqueued) => { // TODO(bt,2025-07-05): Enqueue job token? - info!("Successfully enqueued Artcraft flux pro 1 inpaint. Job token: {}", + info!("Successfully enqueued Artcraft flux pro 1 inpaint. Job token: {}", enqueued.inference_job_token); enqueued.inference_job_token } @@ -115,7 +115,7 @@ pub async fn handle_flux_pro_1_inpaint_artcraft( Ok(TaskEnqueueSuccess { provider: GenerationProvider::Artcraft, - model: Some(GenerationModel::FluxPro1), + model: Some(GenerationModel::FluxProKontextMax), provider_job_id: Some(job_id.to_string()), task_type: TaskType::ImageGeneration, }) diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/flux_pro_1_inpaint/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/flux_pro_1_inpaint/mod.rs index fcdffe4a28..e3040c79e5 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/flux_pro_1_inpaint/mod.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/flux_pro_1_inpaint/mod.rs @@ -1,2 +1,2 @@ pub (super) mod handle_flux_pro_1_inpaint; -pub (crate) mod handle_flux_pro_1_inpaint_artcraft; +mod handle_flux_pro_1_inpaint_artcraft; diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs index e6a4240a08..204905cc17 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs @@ -43,6 +43,7 @@ impl TaskEnqueueSuccess{ Some(GenerationModel::FluxPro1) => Some(TaskModelType::FluxPro1), // NB: This is for inpainting. Some(GenerationModel::FluxPro11) => Some(TaskModelType::FluxPro11), Some(GenerationModel::FluxPro11Ultra) => Some(TaskModelType::FluxPro11Ultra), + Some(GenerationModel::FluxProKontextMax) => Some(TaskModelType::FluxProKontextMax), Some(GenerationModel::GptImage1) => Some(TaskModelType::GptImage1), Some(GenerationModel::Recraft3) => Some(TaskModelType::Recraft3), Some(GenerationModel::Kling21Pro) => Some(TaskModelType::Kling21Pro), diff --git a/crates/desktop/artcraft/src/core/events/generation_events/common.rs b/crates/desktop/artcraft/src/core/events/generation_events/common.rs index 79e8475231..9ada3faea8 100644 --- a/crates/desktop/artcraft/src/core/events/generation_events/common.rs +++ b/crates/desktop/artcraft/src/core/events/generation_events/common.rs @@ -21,6 +21,8 @@ pub enum GenerationModel { FluxPro11, #[serde(rename = "flux_pro_1.1_ultra")] FluxPro11Ultra, + #[serde(rename = "flux_pro_kontext_max")] + FluxProKontextMax, #[serde(rename = "gpt_image_1")] GptImage1, #[serde(rename = "hunyuan_3d_2_0")] diff --git a/crates/desktop/artcraft/src/core/model/contextual_image_edit_models.rs b/crates/desktop/artcraft/src/core/model/contextual_image_edit_models.rs index 5224376a32..06cb2d43ef 100644 --- a/crates/desktop/artcraft/src/core/model/contextual_image_edit_models.rs +++ b/crates/desktop/artcraft/src/core/model/contextual_image_edit_models.rs @@ -7,6 +7,9 @@ use serde_derive::Deserialize; pub enum ContextualImageEditModel { #[serde(rename = "gpt_image_1")] GptImage1, + + #[serde(rename = "flux_pro_kontext_max")] + FluxProKontextMax, // TODO: Flux Kontext models. } diff --git a/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs b/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs index 0de074e517..12e036f3f3 100644 --- a/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs +++ b/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs @@ -22,6 +22,8 @@ pub enum TaskModelType { FluxPro11, #[serde(rename = "flux_pro_1.1_ultra")] FluxPro11Ultra, + #[serde(rename = "flux_pro_kontext_max")] + FluxProKontextMax, #[serde(rename = "gpt_image_1")] GptImage1, #[serde(rename = "recraft_3")] @@ -61,6 +63,7 @@ impl TaskModelType { Self::FluxPro1 => "flux_pro_1", Self::FluxPro11 => "flux_pro_1.1", Self::FluxPro11Ultra => "flux_pro_1.1_ultra", + Self::FluxProKontextMax => "flux_pro_kontext_max", Self::GptImage1 => "gpt_image_1", Self::Recraft3 => "recraft_3", // Video models @@ -83,6 +86,7 @@ impl TaskModelType { "flux_pro_1" => Ok(Self::FluxPro1), "flux_pro_1.1" => Ok(Self::FluxPro11), "flux_pro_1.1_ultra" => Ok(Self::FluxPro11Ultra), + "flux_pro_kontext_max" => Ok(Self::FluxProKontextMax), "gpt_image_1" => Ok(Self::GptImage1), "recraft_3" => Ok(Self::Recraft3), // Video models @@ -108,6 +112,7 @@ impl TaskModelType { Self::FluxPro1, Self::FluxPro11, Self::FluxPro11Ultra, + Self::FluxProKontextMax, Self::GptImage1, Self::Recraft3, // Video models @@ -139,6 +144,7 @@ mod tests { assert_serialization(TaskModelType::FluxPro1, "flux_pro_1"); assert_serialization(TaskModelType::FluxPro11, "flux_pro_1.1"); assert_serialization(TaskModelType::FluxPro11Ultra, "flux_pro_1.1_ultra"); + assert_serialization(TaskModelType::FluxProKontextMax, "flux_pro_kontext_max"); assert_serialization(TaskModelType::GptImage1, "gpt_image_1"); assert_serialization(TaskModelType::Recraft3, "recraft_3"); // Video models @@ -160,6 +166,7 @@ mod tests { assert_eq!(TaskModelType::FluxPro1.to_str(), "flux_pro_1"); assert_eq!(TaskModelType::FluxPro11.to_str(), "flux_pro_1.1"); assert_eq!(TaskModelType::FluxPro11Ultra.to_str(), "flux_pro_1.1_ultra"); + assert_eq!(TaskModelType::FluxProKontextMax.to_str(), "flux_pro_kontext_max"); assert_eq!(TaskModelType::GptImage1.to_str(), "gpt_image_1"); assert_eq!(TaskModelType::Recraft3.to_str(), "recraft_3"); // Video models @@ -181,6 +188,7 @@ mod tests { assert_eq!(TaskModelType::from_str("flux_pro_1").unwrap(), TaskModelType::FluxPro1); assert_eq!(TaskModelType::from_str("flux_pro_1.1").unwrap(), TaskModelType::FluxPro11); assert_eq!(TaskModelType::from_str("flux_pro_1.1_ultra").unwrap(), TaskModelType::FluxPro11Ultra); + assert_eq!(TaskModelType::from_str("flux_pro_kontext_max").unwrap(), TaskModelType::FluxProKontextMax); assert_eq!(TaskModelType::from_str("gpt_image_1").unwrap(), TaskModelType::GptImage1); assert_eq!(TaskModelType::from_str("recraft_3").unwrap(), TaskModelType::Recraft3); // Video models @@ -197,13 +205,14 @@ mod tests { #[test] fn all_variants() { let mut variants = TaskModelType::all_variants(); - assert_eq!(variants.len(), 14); + assert_eq!(variants.len(), 15); // Image models assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Dev)); assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Schnell)); assert_eq!(variants.pop_first(), Some(TaskModelType::FluxPro1)); assert_eq!(variants.pop_first(), Some(TaskModelType::FluxPro11)); assert_eq!(variants.pop_first(), Some(TaskModelType::FluxPro11Ultra)); + assert_eq!(variants.pop_first(), Some(TaskModelType::FluxProKontextMax)); assert_eq!(variants.pop_first(), Some(TaskModelType::GptImage1)); assert_eq!(variants.pop_first(), Some(TaskModelType::Recraft3)); // Video models