From d0c5f769c9afbf7a7afc17b2dcd264eed97e3ca6 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 21 Jun 2025 09:30:14 -0400 Subject: [PATCH] unify how events are handled --- .../image/enqueue_text_to_image_command.rs | 29 ++++++++++++++----- .../enqueue/image/handle_image_artcraft.rs | 18 +++++++----- .../enqueue/image/handle_image_fal.rs | 28 +++++++++++------- .../enqueue/image/handle_image_sora.rs | 10 ------- .../src/core/commands/enqueue/image/mod.rs | 3 +- .../commands/enqueue/image/success_event.rs | 20 +++++++++++++ .../core/events/generation_events/common.rs | 5 ++++ 7 files changed, 78 insertions(+), 35 deletions(-) create mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/image/success_event.rs diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image/enqueue_text_to_image_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image/enqueue_text_to_image_command.rs index f1362b5d3a..87a3cf8294 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image/enqueue_text_to_image_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image/enqueue_text_to_image_command.rs @@ -59,6 +59,10 @@ use storyteller_client::media_files::upload_image_media_file_from_file::upload_i use storyteller_client::utils::api_host::ApiHost; use tauri::{AppHandle, Emitter, Manager, State}; use tempfile::NamedTempFile; +use crate::core::commands::enqueue::image::success_event::SuccessEvent; +use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; +use crate::core::events::generation_events::generation_enqueue_success_event::GenerationEnqueueSuccessEvent; #[derive(Deserialize)] pub struct EnqueueTextToImageRequest { @@ -148,7 +152,17 @@ pub async fn enqueue_text_to_image_command( error_details: None, }) } - Ok(()) => { + Ok(event) => { + let event = GenerationEnqueueSuccessEvent { + action: GenerationAction::GenerateImage, + service: event.service_provider, + model: Some(event.tauri_event_model()), + }; + + if let Err(err) = event.send(&app) { + error!("Failed to emit event: {:?}", err); // Fail open. + } + Ok(EnqueueTextToImageSuccessResponse {}.into()) } } @@ -165,7 +179,7 @@ pub async fn handle_request( fal_task_queue: &FalTaskQueue, sora_creds_manager: &SoraCredentialManager, sora_task_queue: &SoraTaskQueue, -) -> Result<(), InternalImageError> { +) -> Result { match request.model { None => { @@ -173,7 +187,10 @@ pub async fn handle_request( } Some(ImageModel::GptImage1) => { handle_image_sora(&app, request, sora_creds_manager, sora_task_queue).await?; - return Ok(()); + return Ok(SuccessEvent { + service_provider: GenerationServiceProvider::Sora, + model: ImageModel::GptImage1, + }); } _ => { // Fall-through @@ -181,10 +198,8 @@ pub async fn handle_request( }; if fal_creds_manager.has_apparent_api_token()? { - handle_image_fal(&app, request, fal_creds_manager, fal_task_queue).await?; + handle_image_fal(&app, request, fal_creds_manager, fal_task_queue).await } else { - handle_image_artcraft(request, &app, app_env_configs, app_data_root, storyteller_creds_manager).await?; + handle_image_artcraft(request, &app, app_env_configs, app_data_root, storyteller_creds_manager).await } - - Ok(()) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_artcraft.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_artcraft.rs index 281f2289a6..be1a01b91f 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_artcraft.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_artcraft.rs @@ -1,5 +1,6 @@ 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::success_event::SuccessEvent; 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; @@ -42,7 +43,7 @@ pub async fn handle_image_artcraft( app_env_configs: &AppEnvConfigs, app_data_root: &AppDataRoot, storyteller_creds_manager: &StorytellerCredentialManager, -) -> Result<(), InternalImageError> { +) -> Result { let creds = match storyteller_creds_manager.get_credentials()? { Some(creds) => creds, @@ -63,7 +64,7 @@ pub async fn handle_image_artcraft( let uuid_idempotency_token = generate_random_uuid(); - let mut selected_model = None; + let mut selected_model; let job_token = match request.model { None => { @@ -77,7 +78,7 @@ pub async fn handle_image_artcraft( } Some(ImageModel::Flux1Dev) => { info!("enqueue Flux 1 Dev"); - selected_model = Some(GenerationModel::Flux1Dev); + selected_model = ImageModel::Flux1Dev; let request = GenerateFlux1DevTextToImageRequest { uuid_idempotency_token, prompt: request.prompt, @@ -102,7 +103,7 @@ pub async fn handle_image_artcraft( } Some(ImageModel::Flux1Schnell) => { info!("enqueue Flux 1 Schnell"); - selected_model = Some(GenerationModel::Flux1Schnell); + selected_model = ImageModel::Flux1Schnell; let request = GenerateFlux1SchnellTextToImageRequest { uuid_idempotency_token, prompt: request.prompt, @@ -127,7 +128,7 @@ pub async fn handle_image_artcraft( } Some(ImageModel::FluxPro11) => { info!("enqueue Flux Pro 1.1"); - selected_model = Some(GenerationModel::FluxPro11); + selected_model = ImageModel::FluxPro11; let request = GenerateFluxPro11TextToImageRequest { uuid_idempotency_token, prompt: request.prompt, @@ -152,7 +153,7 @@ pub async fn handle_image_artcraft( } Some(ImageModel::FluxPro11Ultra) => { info!("enqueue Flux Pro 1.1 Ultra"); - selected_model = Some(GenerationModel::FluxPro11Ultra); + selected_model = ImageModel::FluxPro11Ultra; let request = GenerateFluxPro11UltraTextToImageRequest { uuid_idempotency_token, prompt: request.prompt, @@ -177,5 +178,8 @@ pub async fn handle_image_artcraft( } }; - Ok(()) + Ok(SuccessEvent { + service_provider: GenerationServiceProvider::Artcraft, + model:selected_model, + }) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_fal.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_fal.rs index 0462e00d01..dda69a7687 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_fal.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_fal.rs @@ -1,5 +1,6 @@ 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::success_event::SuccessEvent; use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent; @@ -23,7 +24,7 @@ pub async fn handle_image_fal( request: EnqueueTextToImageRequest, fal_creds_manager: &FalCredentialManager, fal_task_queue: &FalTaskQueue, -) -> Result<(), InternalImageError> { +) -> Result { let api_key = match fal_creds_manager.get_key()? { Some(key) => key, @@ -48,6 +49,8 @@ pub async fn handle_image_fal( let prompt = request.prompt.as_deref().unwrap_or(""); + let model; + let result = match request.model { None => { return Err(InternalImageError::NoModelSpecified); @@ -63,6 +66,7 @@ pub async fn handle_image_fal( return Err(InternalImageError::AnyhowError(anyhow!("not yet implemented: {:?}", request.model))); } Some(ImageModel::FluxPro11Ultra) => { + model = ImageModel::FluxPro11Ultra; info!("enqueue Flux Pro 1.1 Ultra text-to-image with prompt: {}", prompt); enqueue_flux_pro_11_ultra_text_to_image(FluxPro11UltraTextToImageArgs { prompt, @@ -70,6 +74,7 @@ pub async fn handle_image_fal( }).await } Some(ImageModel::Recraft3) => { + model = ImageModel::Recraft3; info!("enqueue Recraft v3 text-to-image with prompt: {}", prompt); enqueue_recraft3_text_to_image(Recraft3TextToImageArgs { prompt, @@ -82,15 +87,15 @@ pub async fn handle_image_fal( Ok(enqueued) => { info!("Successfully enqueued text to image"); - let event = GenerationEnqueueSuccessEvent { - action: GenerationAction::GenerateImage, - service: GenerationServiceProvider::Fal, - model: None, - }; + //let event = GenerationEnqueueSuccessEvent { + // action: GenerationAction::GenerateImage, + // service: GenerationServiceProvider::Fal, + // model: None, + //}; - if let Err(err) = event.send(app) { - error!("Failed to emit event: {:?}", err); // Fail open. - } + //if let Err(err) = event.send(app) { + // error!("Failed to emit event: {:?}", err); // Fail open. + //} if let Err(err) = fal_task_queue.insert(&enqueued) { error!("Failed to enqueue task: {:?}", err); @@ -115,5 +120,8 @@ pub async fn handle_image_fal( } } - Ok(()) + Ok(SuccessEvent { + service_provider: GenerationServiceProvider::Fal, + model, + }) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_sora.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_sora.rs index 59ad401f1f..8079397854 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_sora.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image/handle_image_sora.rs @@ -90,15 +90,5 @@ pub async fn handle_image_sora( sora_task_queue.insert(&response.task_id)?; - let event = GenerationEnqueueSuccessEvent { - action: GenerationAction::GenerateImage, - service: GenerationServiceProvider::Sora, - model: None, - }; - - if let Err(err) = event.send(app) { - error!("Failed to emit event: {:?}", err); // Fail open. - } - Ok(()) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image/mod.rs index 7e71b73802..1c8332aa3c 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image/mod.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image/mod.rs @@ -2,4 +2,5 @@ pub mod enqueue_text_to_image_command; pub mod handle_image_artcraft; pub mod handle_image_fal; pub mod handle_image_sora; -pub mod internal_image_error; \ No newline at end of file +pub mod internal_image_error; +mod success_event; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image/success_event.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image/success_event.rs new file mode 100644 index 0000000000..a0688ca693 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image/success_event.rs @@ -0,0 +1,20 @@ +use crate::core::events::generation_events::common::{GenerationModel, GenerationServiceProvider}; +use crate::core::model::image_models::ImageModel; + +pub struct SuccessEvent { + pub service_provider: GenerationServiceProvider, + pub model: ImageModel, +} + +impl SuccessEvent { + pub fn tauri_event_model(&self) -> GenerationModel { + match self.model { + ImageModel::Flux1Dev => GenerationModel::Flux1Dev, + ImageModel::Flux1Schnell => GenerationModel::Flux1Schnell, + ImageModel::FluxPro11 => GenerationModel::FluxPro11, + ImageModel::FluxPro11Ultra => GenerationModel::FluxPro11Ultra, + ImageModel::GptImage1 => GenerationModel::GptImage1, + ImageModel::Recraft3 => GenerationModel::Recraft3, + } + } +} \ No newline at end of file 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 df7d8dfe04..9f766bd61d 100644 --- a/crates/desktop/artcraft/src/core/events/generation_events/common.rs +++ b/crates/desktop/artcraft/src/core/events/generation_events/common.rs @@ -19,6 +19,8 @@ pub enum GenerationModel { FluxPro11, #[serde(rename = "flux_pro_1.1_ultra")] FluxPro11Ultra, + #[serde(rename = "gpt_image_1")] + GptImage1, #[serde(rename = "kling_1.6")] Kling1_6, #[serde(rename = "kling_2.0")] @@ -27,6 +29,9 @@ pub enum GenerationModel { Kling21Master, #[serde(rename = "kling_2.1_pro")] Kling21Pro, + #[serde(rename = "recraft_3")] + Recraft3, + #[serde(rename = "sora")] Sora, #[serde(rename = "seedance_1.0_lite")] Seedance10Lite,